<commit_before>//! The SMTP transport sends emails using the SMTP protocol.
//!
//! This SMTP client follows [RFC
//! 5321](https://tools.ietf.org/html/rfc5321), and is designed to efficiently send emails from an
//! application to a relay email server, as it relies as much as possible on the relay server
//! for sanity and RFC compliance checks.
//!
//! It implements the following extensions:
//!
//! * 8BITMIME ([RFC 6152](https://tools.ietf.org/html/rfc6152))
//! * AUTH ([RFC 4954](http://tools.ietf.org/html/rfc4954)) with PLAIN, LOGIN and XOAUTH2 mechanisms
//! * STARTTLS ([RFC 2487](http://tools.ietf.org/html/rfc2487))
//! * SMTPUTF8 ([RFC 6531](http://tools.ietf.org/html/rfc6531))
//!

use native_tls::TlsConnector;
use smtp::authentication::{Credentials, Mechanism, DEFAULT_ENCRYPTED_MECHANISMS,
                           DEFAULT_UNENCRYPTED_MECHANISMS};
use smtp::client::net::ClientTlsParameters;
use smtp::client::net::DEFAULT_TLS_PROTOCOLS;
use smtp::client::InnerClient;
use smtp::commands::*;
use smtp::error::{Error, SmtpResult};
use smtp::extension::{ClientId, Extension, MailBodyParameter, MailParameter, ServerInfo};
use std::net::{SocketAddr, ToSocketAddrs};
use std::time::Duration;
use {SendableEmail, Transport};

pub mod authentication;
pub mod client;
pub mod commands;
pub mod error;
pub mod extension;
pub mod response;
pub mod util;

// Registered port numbers:
// https://www.iana.
// org/assignments/service-names-port-numbers/service-names-port-numbers.xhtml

/// Default smtp port
pub const SMTP_PORT: u16 = 25;

/// Default submission port
pub const SUBMISSION_PORT: u16 = 587;

/// How to apply TLS to a client connection
#[derive(Clone)]
#[allow(missing_debug_implementations)]
pub enum ClientSecurity {
    /// Insecure connection
    None,
    /// Use `STARTTLS` when available
    Opportunistic(ClientTlsParameters),
    /// Always use `STARTTLS`
    Required(ClientTlsParameters),
    /// Use TLS wrapped connection without negotiation
    /// Non RFC-compliant, should only be used if the server does not support STARTTLS.
    Wrapper(ClientTlsParameters),
}

/// Configures connection reuse behavior
#[derive(Clone, Debug, Copy)]
#[cfg_attr(feature = "serde-impls", derive(Serialize, Deserialize))]
pub enum ConnectionReuseParameters {
    /// Unlimited connection reuse
    ReuseUnlimited,
    /// Maximum number of connection reuse
    ReuseLimited(u16),
    /// Disable connection reuse, close connection after each transaction
    NoReuse,
}

/// Contains client configuration
#[allow(missing_debug_implementations)]
pub struct SmtpClient {
    /// Enable connection reuse
    connection_reuse: ConnectionReuseParameters,
    /// Name sent during EHLO
    hello_name: ClientId,
    /// Credentials
    credentials: Option<Credentials>,
    /// Socket we are connecting to
    server_addr: SocketAddr,
    /// TLS security configuration
    security: ClientSecurity,
    /// Enable UTF8 mailboxes in envelope or headers
    smtp_utf8: bool,
    /// Optional enforced authentication mechanism
    authentication_mechanism: Option<Mechanism>,
    /// Define network timeout
    /// It can be changed later for specific needs (like a different timeout for each SMTP command)
    timeout: Option<Duration>,
}

/// Builder for the SMTP `SmtpTransport`
impl SmtpClient {
    /// Creates a new SMTP client
    ///
    /// Defaults are:
    ///
    /// * No connection reuse
    /// * No authentication
    /// * No SMTPUTF8 support
    /// * A 60 seconds timeout for smtp commands
    pub fn new<A: ToSocketAddrs>(addr: A, security: ClientSecurity) -> Result<SmtpClient, Error> {
        let mut addresses = addr.to_socket_addrs()?;

        match addresses.next() {
            Some(addr) => Ok(SmtpClient {
                server_addr: addr,
                security,
                smtp_utf8: false,
                credentials: None,
                connection_reuse: ConnectionReuseParameters::NoReuse,
                hello_name: ClientId::hostname(),
                authentication_mechanism: None,
                timeout: Some(Duration::new(60, 0)),
            }),
            None => Err(Error::Resolution),
        }
    }

    /// Simple and secure transport, should be used when possible.
    /// Creates an encrypted transport over submission port, using the provided domain
    /// to validate TLS certificates.
    pub fn new_simple(domain: &str) -> Result<SmtpClient, Error> {
        let mut tls_builder = TlsConnector::builder()?;
        tls_builder.supported_protocols(DEFAULT_TLS_PROTOCOLS)?;

        let tls_parameters =
            ClientTlsParameters::new(domain.to_string(), tls_builder.build().unwrap());

        SmtpClient::new(
            (domain, SUBMISSION_PORT),
            ClientSecurity::Required(tls_parameters),
        )
    }

    /// Creates a new local SMTP client to port 25
    pub fn new_unencrypted_localhost() -> Result<SmtpClient, Error> {
        SmtpClient::new(("localhost", SMTP_PORT), ClientSecurity::None)
    }

    /// Enable SMTPUTF8 if the server supports it
    pub fn smtp_utf8(mut self, enabled: bool) -> SmtpClient {
        self.smtp_utf8 = enabled;
        self
    }

    /// Set the name used during EHLO
    pub fn hello_name(mut self, name: ClientId) -> SmtpClient {
        self.hello_name = name;
        self
    }

    /// Enable connection reuse
    pub fn connection_reuse(mut self, parameters: ConnectionReuseParameters) -> SmtpClient {
        self.connection_reuse = parameters;
        self
    }

    /// Set the client credentials
    pub fn credentials<S: Into<Credentials>>(mut self, credentials: S) -> SmtpClient {
        self.credentials = Some(credentials.into());
        self
    }

    /// Set the authentication mechanism to use
    pub fn authentication_mechanism(mut self, mechanism: Mechanism) -> SmtpClient {
        self.authentication_mechanism = Some(mechanism);
        self
    }

    /// Set the timeout duration
    pub fn timeout(mut self, timeout: Option<Duration>) -> SmtpClient {
        self.timeout = timeout;
        self
    }

    /// Build the SMTP client
    ///
    /// It does not connect to the server, but only creates the `SmtpTransport`
    pub fn transport(self) -> SmtpTransport {
        SmtpTransport::new(self)
    }
}

/// Represents the state of a client
#[derive(Debug)]
struct State {
    /// Panic state
    pub panic: bool,
    /// Connection reuse counter
    pub connection_reuse_count: u16,
}

/// Structure that implements the high level SMTP client
#[allow(missing_debug_implementations)]
pub struct SmtpTransport {
    /// Information about the server
    /// Value is None before HELO/EHLO
    server_info: Option<ServerInfo>,
    /// SmtpTransport variable states
    state: State,
    /// Information about the client
    client_info: SmtpClient,
    /// Low level client
    client: InnerClient,
}

macro_rules! try_smtp (
    ($err: expr, $client: ident) => ({
        match $err {
            Ok(val) => val,
            Err(err) => {
                if !$client.state.panic {
                    $client.state.panic = true;
                    $client.close();
                }
                return Err(From::from(err))
            },
        }
    })
);

impl<'a> SmtpTransport {
    /// Creates a new SMTP client
    ///
    /// It does not connect to the server, but only creates the `SmtpTransport`
    pub fn new(builder: SmtpClient) -> SmtpTransport {
        let client = InnerClient::new();

        SmtpTransport {
            client,
            server_info: None,
            client_info: builder,
            state: State {
                panic: false,
                connection_reuse_count: 0,
            },
        }
    }

    /// Gets the EHLO response and updates server information
    fn ehlo(&mut self) -> SmtpResult {
        // Extended Hello
        let ehlo_response = try_smtp!(
            self.client.command(EhloCommand::new(ClientId::new(
                self.client_info.hello_name.to_string()
            ),)),
            self
        );

        self.server_info = Some(try_smtp!(ServerInfo::from_response(&ehlo_response), self));

        // Print server information
        debug!("server {}", self.server_info.as_ref().unwrap());

        Ok(ehlo_response)
    }

    /// Reset the client state
    pub fn close(&mut self) {
        // Close the SMTP transaction if needed
        self.client.close();

        // Reset the client state
        self.server_info = None;
        self.state.panic = false;
        self.state.connection_reuse_count = 0;
    }
}

impl<'a> Transport<'a> for SmtpTransport {
    type Result = SmtpResult;

    /// Sends an email
    #[cfg_attr(feature = "cargo-clippy", allow(match_same_arms, cyclomatic_complexity))]
    fn send(&mut self, email: SendableEmail) -> SmtpResult {
        let message_id = email.message_id().to_string();

        // Check if the connection is still available
        if (self.state.connection_reuse_count > 0) && (!self.client.is_connected()) {
            self.close();
        }

        if self.state.connection_reuse_count == 0 {
            self.client.connect(
                &self.client_info.server_addr,
                match self.client_info.security {
                    ClientSecurity::Wrapper(ref tls_parameters) => Some(tls_parameters),
                    _ => None,
                },
            )?;

            self.client.set_timeout(self.client_info.timeout)?;

            // Log the connection
            info!("connection established to {}", self.client_info.server_addr);

            self.ehlo()?;

            match (
                &self.client_info.security.clone(),
                self.server_info
                    .as_ref()
                    .unwrap()
                    .supports_feature(Extension::StartTls),
            ) {
                (&ClientSecurity::Required(_), false) => {
                    return Err(From::from("Could not encrypt connection, aborting"))
                }
                (&ClientSecurity::Opportunistic(_), false) => (),
                (&ClientSecurity::None, _) => (),
                (&ClientSecurity::Wrapper(_), _) => (),
                (&ClientSecurity::Opportunistic(ref tls_parameters), true)
                | (&ClientSecurity::Required(ref tls_parameters), true) => {
                    try_smtp!(self.client.command(StarttlsCommand), self);
                    try_smtp!(self.client.upgrade_tls_stream(tls_parameters), self);

                    debug!("connection encrypted");

                    // Send EHLO again
                    self.ehlo()?;
                }
            }

            if self.client_info.credentials.is_some() {
                let mut found = false;

                // Compute accepted mechanism
                let accepted_mechanisms = match self.client_info.authentication_mechanism {
                    Some(mechanism) => vec![mechanism],
                    None => {
                        if self.client.is_encrypted() {
                            DEFAULT_ENCRYPTED_MECHANISMS.to_vec()
                        } else {
                            DEFAULT_UNENCRYPTED_MECHANISMS.to_vec()
                        }
                    }
                };

                for mechanism in accepted_mechanisms {
                    if self.server_info
                        .as_ref()
                        .unwrap()
                        .supports_auth_mechanism(mechanism)
                    {
                        found = true;
                        try_smtp!(
                            self.client
                                .auth(mechanism, self.client_info.credentials.as_ref().unwrap(),),
                            self
                        );
                        break;
                    }
                }

                if !found {
                    info!("No supported authentication mechanisms available");
                }
            }
        }

        // Mail
        let mut mail_options = vec![];

        if self.server_info
            .as_ref()
            .unwrap()
            .supports_feature(Extension::EightBitMime)
        {
            mail_options.push(MailParameter::Body(MailBodyParameter::EightBitMime));
        }

        if self.server_info
            .as_ref()
            .unwrap()
            .supports_feature(Extension::SmtpUtfEight) && self.client_info.smtp_utf8
        {
            mail_options.push(MailParameter::SmtpUtfEight);
        }

        try_smtp!(
            self.client.command(MailCommand::new(
                email.envelope().from().cloned(),
                mail_options,
            )),
            self
        );

        // Log the mail command
        info!(
            "{}: from=<{}>",
            message_id,
            match email.envelope().from() {
                Some(address) => address.to_string(),
                None => "".to_string(),
            }
        );

        // Recipient
        for to_address in email.envelope().to() {
            try_smtp!(
                self.client
                    .command(RcptCommand::new(to_address.clone(), vec![])),
                self
            );
            // Log the rcpt command
            info!("{}: to=<{}>", message_id, to_address);
        }

        // Data
        try_smtp!(self.client.command(DataCommand), self);

        // Message content
        let result = self.client.message(Box::new(email.message()));

        if result.is_ok() {
            // Increment the connection reuse counter
            self.state.connection_reuse_count += 1;

            // Log the message
            info!(
                "{}: conn_use={}, status=sent ({})",
                message_id,
                self.state.connection_reuse_count,
                result
                    .as_ref()
                    .ok()
                    .unwrap()
                    .message
                    .iter()
                    .next()
                    .unwrap_or(&"no response".to_string())
            );
        }

        // Test if we can reuse the existing connection
        match self.client_info.connection_reuse {
            ConnectionReuseParameters::ReuseLimited(limit)
                if self.state.connection_reuse_count >= limit =>
            {
                self.close()
            }
            ConnectionReuseParameters::NoReuse => self.close(),
            _ => (),
        }

        result
    }
}
<commit_msg>feat(transport): Use submissions port by default<commit_after>//! The SMTP transport sends emails using the SMTP protocol.
//!
//! This SMTP client follows [RFC
//! 5321](https://tools.ietf.org/html/rfc5321), and is designed to efficiently send emails from an
//! application to a relay email server, as it relies as much as possible on the relay server
//! for sanity and RFC compliance checks.
//!
//! It implements the following extensions:
//!
//! * 8BITMIME ([RFC 6152](https://tools.ietf.org/html/rfc6152))
//! * AUTH ([RFC 4954](http://tools.ietf.org/html/rfc4954)) with PLAIN, LOGIN and XOAUTH2 mechanisms
//! * STARTTLS ([RFC 2487](http://tools.ietf.org/html/rfc2487))
//! * SMTPUTF8 ([RFC 6531](http://tools.ietf.org/html/rfc6531))
//!

use native_tls::TlsConnector;
use smtp::authentication::{Credentials, Mechanism, DEFAULT_ENCRYPTED_MECHANISMS,
                           DEFAULT_UNENCRYPTED_MECHANISMS};
use smtp::client::net::ClientTlsParameters;
use smtp::client::net::DEFAULT_TLS_PROTOCOLS;
use smtp::client::InnerClient;
use smtp::commands::*;
use smtp::error::{Error, SmtpResult};
use smtp::extension::{ClientId, Extension, MailBodyParameter, MailParameter, ServerInfo};
use std::net::{SocketAddr, ToSocketAddrs};
use std::time::Duration;
use {SendableEmail, Transport};

pub mod authentication;
pub mod client;
pub mod commands;
pub mod error;
pub mod extension;
pub mod response;
pub mod util;

// Registered port numbers:
// https://www.iana.
// org/assignments/service-names-port-numbers/service-names-port-numbers.xhtml

/// Default smtp port
pub const SMTP_PORT: u16 = 25;
/// Default submission port
pub const SUBMISSION_PORT: u16 = 587;
/// Default submission over TLS port
pub const SUBMISSIONS_PORT: u16 = 465;

/// How to apply TLS to a client connection
#[derive(Clone)]
#[allow(missing_debug_implementations)]
pub enum ClientSecurity {
    /// Insecure connection only (for testing purposes)
    None,
    /// Start with insecure connection and use `STARTTLS` when available
    Opportunistic(ClientTlsParameters),
    /// Start with insecure connection and require `STARTTLS`
    Required(ClientTlsParameters),
    /// Use TLS wrapped connection
    Wrapper(ClientTlsParameters),
}

/// Configures connection reuse behavior
#[derive(Clone, Debug, Copy)]
#[cfg_attr(feature = "serde-impls", derive(Serialize, Deserialize))]
pub enum ConnectionReuseParameters {
    /// Unlimited connection reuse
    ReuseUnlimited,
    /// Maximum number of connection reuse
    ReuseLimited(u16),
    /// Disable connection reuse, close connection after each transaction
    NoReuse,
}

/// Contains client configuration
#[allow(missing_debug_implementations)]
pub struct SmtpClient {
    /// Enable connection reuse
    connection_reuse: ConnectionReuseParameters,
    /// Name sent during EHLO
    hello_name: ClientId,
    /// Credentials
    credentials: Option<Credentials>,
    /// Socket we are connecting to
    server_addr: SocketAddr,
    /// TLS security configuration
    security: ClientSecurity,
    /// Enable UTF8 mailboxes in envelope or headers
    smtp_utf8: bool,
    /// Optional enforced authentication mechanism
    authentication_mechanism: Option<Mechanism>,
    /// Define network timeout
    /// It can be changed later for specific needs (like a different timeout for each SMTP command)
    timeout: Option<Duration>,
}

/// Builder for the SMTP `SmtpTransport`
impl SmtpClient {
    /// Creates a new SMTP client
    ///
    /// Defaults are:
    ///
    /// * No connection reuse
    /// * No authentication
    /// * No SMTPUTF8 support
    /// * A 60 seconds timeout for smtp commands
    pub fn new<A: ToSocketAddrs>(addr: A, security: ClientSecurity) -> Result<SmtpClient, Error> {
        let mut addresses = addr.to_socket_addrs()?;

        match addresses.next() {
            Some(addr) => Ok(SmtpClient {
                server_addr: addr,
                security,
                smtp_utf8: false,
                credentials: None,
                connection_reuse: ConnectionReuseParameters::NoReuse,
                hello_name: ClientId::hostname(),
                authentication_mechanism: None,
                timeout: Some(Duration::new(60, 0)),
            }),
            None => Err(Error::Resolution),
        }
    }

    /// Simple and secure transport, should be used when possible.
    /// Creates an encrypted transport over submissions port, using the provided domain
    /// to validate TLS certificates.
    pub fn new_simple(domain: &str) -> Result<SmtpClient, Error> {
        let mut tls_builder = TlsConnector::builder()?;
        tls_builder.supported_protocols(DEFAULT_TLS_PROTOCOLS)?;

        let tls_parameters =
            ClientTlsParameters::new(domain.to_string(), tls_builder.build().unwrap());

        SmtpClient::new(
            (domain, SUBMISSIONS_PORT),
            ClientSecurity::Wrapper(tls_parameters),
        )
    }

    /// Creates a new local SMTP client to port 25
    pub fn new_unencrypted_localhost() -> Result<SmtpClient, Error> {
        SmtpClient::new(("localhost", SMTP_PORT), ClientSecurity::None)
    }

    /// Enable SMTPUTF8 if the server supports it
    pub fn smtp_utf8(mut self, enabled: bool) -> SmtpClient {
        self.smtp_utf8 = enabled;
        self
    }

    /// Set the name used during EHLO
    pub fn hello_name(mut self, name: ClientId) -> SmtpClient {
        self.hello_name = name;
        self
    }

    /// Enable connection reuse
    pub fn connection_reuse(mut self, parameters: ConnectionReuseParameters) -> SmtpClient {
        self.connection_reuse = parameters;
        self
    }

    /// Set the client credentials
    pub fn credentials<S: Into<Credentials>>(mut self, credentials: S) -> SmtpClient {
        self.credentials = Some(credentials.into());
        self
    }

    /// Set the authentication mechanism to use
    pub fn authentication_mechanism(mut self, mechanism: Mechanism) -> SmtpClient {
        self.authentication_mechanism = Some(mechanism);
        self
    }

    /// Set the timeout duration
    pub fn timeout(mut self, timeout: Option<Duration>) -> SmtpClient {
        self.timeout = timeout;
        self
    }

    /// Build the SMTP client
    ///
    /// It does not connect to the server, but only creates the `SmtpTransport`
    pub fn transport(self) -> SmtpTransport {
        SmtpTransport::new(self)
    }
}

/// Represents the state of a client
#[derive(Debug)]
struct State {
    /// Panic state
    pub panic: bool,
    /// Connection reuse counter
    pub connection_reuse_count: u16,
}

/// Structure that implements the high level SMTP client
#[allow(missing_debug_implementations)]
pub struct SmtpTransport {
    /// Information about the server
    /// Value is None before HELO/EHLO
    server_info: Option<ServerInfo>,
    /// SmtpTransport variable states
    state: State,
    /// Information about the client
    client_info: SmtpClient,
    /// Low level client
    client: InnerClient,
}

macro_rules! try_smtp (
    ($err: expr, $client: ident) => ({
        match $err {
            Ok(val) => val,
            Err(err) => {
                if !$client.state.panic {
                    $client.state.panic = true;
                    $client.close();
                }
                return Err(From::from(err))
            },
        }
    })
);

impl<'a> SmtpTransport {
    /// Creates a new SMTP client
    ///
    /// It does not connect to the server, but only creates the `SmtpTransport`
    pub fn new(builder: SmtpClient) -> SmtpTransport {
        let client = InnerClient::new();

        SmtpTransport {
            client,
            server_info: None,
            client_info: builder,
            state: State {
                panic: false,
                connection_reuse_count: 0,
            },
        }
    }

    /// Gets the EHLO response and updates server information
    fn ehlo(&mut self) -> SmtpResult {
        // Extended Hello
        let ehlo_response = try_smtp!(
            self.client.command(EhloCommand::new(ClientId::new(
                self.client_info.hello_name.to_string()
            ),)),
            self
        );

        self.server_info = Some(try_smtp!(ServerInfo::from_response(&ehlo_response), self));

        // Print server information
        debug!("server {}", self.server_info.as_ref().unwrap());

        Ok(ehlo_response)
    }

    /// Reset the client state
    pub fn close(&mut self) {
        // Close the SMTP transaction if needed
        self.client.close();

        // Reset the client state
        self.server_info = None;
        self.state.panic = false;
        self.state.connection_reuse_count = 0;
    }
}

impl<'a> Transport<'a> for SmtpTransport {
    type Result = SmtpResult;

    /// Sends an email
    #[cfg_attr(feature = "cargo-clippy", allow(match_same_arms, cyclomatic_complexity))]
    fn send(&mut self, email: SendableEmail) -> SmtpResult {
        let message_id = email.message_id().to_string();

        // Check if the connection is still available
        if (self.state.connection_reuse_count > 0) && (!self.client.is_connected()) {
            self.close();
        }

        if self.state.connection_reuse_count == 0 {
            self.client.connect(
                &self.client_info.server_addr,
                match self.client_info.security {
                    ClientSecurity::Wrapper(ref tls_parameters) => Some(tls_parameters),
                    _ => None,
                },
            )?;

            self.client.set_timeout(self.client_info.timeout)?;

            // Log the connection
            info!("connection established to {}", self.client_info.server_addr);

            self.ehlo()?;

            match (
                &self.client_info.security.clone(),
                self.server_info
                    .as_ref()
                    .unwrap()
                    .supports_feature(Extension::StartTls),
            ) {
                (&ClientSecurity::Required(_), false) => {
                    return Err(From::from("Could not encrypt connection, aborting"))
                }
                (&ClientSecurity::Opportunistic(_), false) => (),
                (&ClientSecurity::None, _) => (),
                (&ClientSecurity::Wrapper(_), _) => (),
                (&ClientSecurity::Opportunistic(ref tls_parameters), true)
                | (&ClientSecurity::Required(ref tls_parameters), true) => {
                    try_smtp!(self.client.command(StarttlsCommand), self);
                    try_smtp!(self.client.upgrade_tls_stream(tls_parameters), self);

                    debug!("connection encrypted");

                    // Send EHLO again
                    self.ehlo()?;
                }
            }

            if self.client_info.credentials.is_some() {
                let mut found = false;

                // Compute accepted mechanism
                let accepted_mechanisms = match self.client_info.authentication_mechanism {
                    Some(mechanism) => vec![mechanism],
                    None => {
                        if self.client.is_encrypted() {
                            DEFAULT_ENCRYPTED_MECHANISMS.to_vec()
                        } else {
                            DEFAULT_UNENCRYPTED_MECHANISMS.to_vec()
                        }
                    }
                };

                for mechanism in accepted_mechanisms {
                    if self.server_info
                        .as_ref()
                        .unwrap()
                        .supports_auth_mechanism(mechanism)
                    {
                        found = true;
                        try_smtp!(
                            self.client
                                .auth(mechanism, self.client_info.credentials.as_ref().unwrap(),),
                            self
                        );
                        break;
                    }
                }

                if !found {
                    info!("No supported authentication mechanisms available");
                }
            }
        }

        // Mail
        let mut mail_options = vec![];

        if self.server_info
            .as_ref()
            .unwrap()
            .supports_feature(Extension::EightBitMime)
        {
            mail_options.push(MailParameter::Body(MailBodyParameter::EightBitMime));
        }

        if self.server_info
            .as_ref()
            .unwrap()
            .supports_feature(Extension::SmtpUtfEight) && self.client_info.smtp_utf8
        {
            mail_options.push(MailParameter::SmtpUtfEight);
        }

        try_smtp!(
            self.client.command(MailCommand::new(
                email.envelope().from().cloned(),
                mail_options,
            )),
            self
        );

        // Log the mail command
        info!(
            "{}: from=<{}>",
            message_id,
            match email.envelope().from() {
                Some(address) => address.to_string(),
                None => "".to_string(),
            }
        );

        // Recipient
        for to_address in email.envelope().to() {
            try_smtp!(
                self.client
                    .command(RcptCommand::new(to_address.clone(), vec![])),
                self
            );
            // Log the rcpt command
            info!("{}: to=<{}>", message_id, to_address);
        }

        // Data
        try_smtp!(self.client.command(DataCommand), self);

        // Message content
        let result = self.client.message(Box::new(email.message()));

        if result.is_ok() {
            // Increment the connection reuse counter
            self.state.connection_reuse_count += 1;

            // Log the message
            info!(
                "{}: conn_use={}, status=sent ({})",
                message_id,
                self.state.connection_reuse_count,
                result
                    .as_ref()
                    .ok()
                    .unwrap()
                    .message
                    .iter()
                    .next()
                    .unwrap_or(&"no response".to_string())
            );
        }

        // Test if we can reuse the existing connection
        match self.client_info.connection_reuse {
            ConnectionReuseParameters::ReuseLimited(limit)
                if self.state.connection_reuse_count >= limit =>
            {
                self.close()
            }
            ConnectionReuseParameters::NoReuse => self.close(),
            _ => (),
        }

        result
    }
}
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! # SOS VGA
//! Code for interacting with the system's VGA buffer.
//!
//! # Features
//! This crate exposes the following feature flags:
//! + `system_term`: Use the VGA buffer as the system terminal. This provides
//!   implementations of the `print!` and `println!` macros.
//! + `log`: provides a `log!` macro for kernel log messages.
#![crate_name = "sos_vga"]
#![crate_type = "lib"]
#![feature( const_fn
          , slice_patterns
          , unique
          , question_mark )]
#![cfg_attr(feature = "system_term", feature(lang_items))]
#![cfg_attr(feature = "clippy", feature(plugin))]
#![cfg_attr(feature = "clippy", plugin(clippy))]
#![no_std]

#[cfg(feature = "system_term")]
extern crate spin;

use core::{mem, ptr};
use core::fmt::{Write, Result};
// use core::ptr::Unique;

#[cfg(feature = "system_term")]
use spin::Mutex;

/// The system's global VGA terminal
#[cfg(feature = "system_term")]
pub static CONSOLE: Mutex<Terminal>
    = Mutex::new(unsafe { Terminal::new(
         Palette::new(Color::LightGrey, Color::Black )
       , 0xB8000
    )});

// TODO: support varying VGA screen sizes?
pub const X_MAX: usize = 80;
pub const Y_MAX: usize = 25;

/// The type signature fot the actual VGA buffer
pub type Buffer = [[Char; X_MAX]; Y_MAX];

const ANSI_ESCAPE: &'static str = "\x1b";
const FG_MASK: u8 = 0b0000_1111;
const BG_MASK: u8 = 0b1111_0000;

/// VGA color codes
#[derive(Eq, PartialEq, Ord, PartialOrd, Copy, Clone)]
#[repr(u8)]
pub enum Color { Black = 0
               , Blue = 1
               , Green = 2
               , Cyan = 3
               , Red = 4
               , Magenta = 5
               , Brown = 6
               , LightGrey = 7
               , DarkGrey = 8
               , LightBlue = 9
               , LightGreen = 10
               , LightCyan = 11
               , LightRed = 12
               , LightMagenta = 13
               , Yellow = 14
               , White = 15
               }

#[derive(Eq, PartialEq, Ord, PartialOrd, Copy, Clone)]
pub struct Palette(u8);

impl Palette {
    /// Returns a `Palette` with the given foreground and background color.
    pub const fn new(fg: Color, bg: Color) -> Self {
        Palette( (bg as u8) << 4 | (fg as u8) )
    }

    /// Returns a new `Palette` with this palette's background color, and
    /// the specified foreground color.
    pub fn set_foreground(&self, fg: Color) -> Self {
        Palette( ( self.0 & BG_MASK) | (fg as u8 & FG_MASK) )
    }

    /// Returns a new `Palette` with this palette's foreground color, and
    /// the specified background color.
    pub fn set_background(&self, bg: Color) -> Self {
        Palette( ( (bg as u8) << 4 & BG_MASK) | (self.0 & FG_MASK) )
    }

    /// Returns this `Palette`'s foreground color.
    pub fn foreground(&self) -> Color {
        unsafe { mem::transmute(self.0 & FG_MASK) }
    }

    /// Returns this `Palette`'s background color.
    pub fn background(&self) -> Color {
        unsafe { mem::transmute((self.0 & BG_MASK) >> 4) }
    }

}
/// A colored VGA character.
#[derive(Copy, Clone)]
#[repr(C)]
pub struct Char { pub ascii: u8
                , pub colors: Palette
                }

impl Char {
    pub const fn empty() -> Self {
        Char { ascii: 0, colors: Palette(0) }
    }
}

pub struct Terminal { buffer: ptr::Unique<Buffer>
                    , x: usize
                    , y: usize
                    , colors: Palette
                    }
impl Terminal {

    /// Constructs a new `Terminal` for abuffer starting at the given address.
    ///
    /// # Arguments:
    /// + `colors`: the default color palette for the terminal
    /// + `buffer_start`: the address of the to the memory location where
    ///                   the terminal's VGA buffer begins
    ///
    /// # Unsafe due to:
    /// + Casting a raw address to an array
    pub const unsafe fn new( colors: Palette
                           , buffer_start: usize)
                           -> Terminal {
        Terminal { x: 0, y: 0
                 , colors: colors
                 , buffer: ptr::Unique::new(buffer_start as *mut _)
                 }
    }

    #[inline]
    fn buffer(&mut self) -> &mut  Buffer {
        unsafe { self.buffer.get_mut() }
    }

    /// Set the color palette used for writing subsequent characters.
    pub fn set_colors(&mut self, bg: Color, fg: Color) -> &mut Self {
        self.colors = Palette::new(bg,fg);
        self
    }

    /// Scrolls the terminal one row.
    fn scroll(&mut self) {
        // // construct an iterator over the whole buffer
        // let mut rows = self.buffer()
        //                    .iter_mut();
        //
        // // the current row in the buffer
        // let mut current = rows.next()
        //                    .unwrap();
        //
        // while let Some(next) = rows.next() {
        //     // while there are rows remaining in the iterator, swap the
        //     // next row with the current one (moving it back by one)
        //     mem::swap(current, next);
        //     // and advance our pointer to the current row.
        //     current = next;
        // }
        let buffer: &mut _ = self.buffer();
        for row in 1..Y_MAX {
            for col in 0 .. X_MAX {
                unsafe {
                    let last = ptr::read_volatile(&buffer[row][col]);
                    ptr::write_volatile( &mut buffer[row - 1][col], last);
                }
            }
        }
        // empty the last line in the buffer
        unsafe {
            for col in 0..X_MAX {
                ptr::write_volatile(&mut buffer[Y_MAX - 1][col], Char::empty());
            }
        }
    }

    /// Clear the terminal
    pub fn clear(&mut self) -> &mut Self {
        // to clear the terminal, we just zero out the whole buffer.
        {
            let buffer: &mut _ =self.buffer();
            for row in 0..Y_MAX {
                for col in 0 .. X_MAX {
                    unsafe {
                        ptr::write_volatile(
                            &mut buffer[row][col], Char::empty()
                        );
                    }
                }
            }
        }

        self.x = 0;
        self.y = 0;
        self
    }

    /// Write the given byte to the terminal, and advance the cursor position.
    pub fn write_byte(&mut self, byte: u8) -> &mut Self {
        if byte == b'\n' {
            // if the byte is a newline, we just advance to the next line
            // and reset the column position.
            self.x = 0;
            self.y += 1;
        } else {
            // otherwise, it's a regular character, so we just set the
            // byte at the current position in the buffer to that
            // character (with the current color palette)
            //
            let y = self.y;
            let x = self.x;
            unsafe {
                ptr::write_volatile( &mut (self.buffer()[y][x])
                                   , Char { ascii: byte, colors: self.colors }
                                   );
            }
            // and advance our column position by one
            self.x += 1;

            if self.x >= X_MAX {
                // if we've reached the end of the line, advance to the next
                self.x = 0;
                self.y += 1;
            }
        }
        if self.y >= Y_MAX {
            // if we've reached the bottom of the terminal, scroll.
            self.scroll();
            self.y -= 1;
        }
        self
    }

    fn handle_ansi_escape(&self, escape_code: &str) -> Result {
        match escape_code.as_bytes() {
            // `\x1b[3Nm` sets the foreground color to N.
            &[0x1b, b'[', b'3', n, b'm'] => {
                unsafe { self.colors
                             .set_foreground(mem::transmute(n - 48)); }
                Ok(())
            }
            // `\x1b[4Nm` sets the background color to N
          , &[0x1b, b'[', b'4', n, b'm'] => {
                unsafe { self.colors
                             .set_background(mem::transmute(n - 48)); }
                Ok(())
            }
          , _ => unimplemented!()
        }
        // let escape_seq: &str = bytes.take_while(|b| b != b'm')
        //                       .collect::<&str>();
        // match escape_seq {
        //     [b'3', n] => unsafe {
        //         self.colors.set_foreground(mem::transmute(n - 48))
        //     }
        // }
        // while let Some(byte) = bytes.next() {
        //     match *byte {
        //         // we've recieved an ANSI escape sequence.
        //         // this basically enters a mediocre FSM for matching ANSI
        //         // control codes.
        //         0x1b => match *next_ansi_byte!(bytes) {
        //             // handle multi-char ANSI escapes
        //             b'[' => match *next_ansi_byte!(bytes) {
        //                 // foreground color code
        //                 fg @ 30 ... 37 => {
        //                     if !(*next_ansi_byte!(bytes) == b'm') {
        //                         unsafe {
        //                             let color: vga::Color
        //                                 = mem::transmute(fg - 30);
        //                             self.colors
        //                                 .set_foreground(color);
        //                         }
        //
        //                     }
        //                 }
        //                 // background color code
        //               , 40 ... 47 => {
        //
        //                 }
        //               , _ => unimplemented!()
        //             }
        //           , _    => unimplemented!()
        //         }
        //         // otherwise, treat the byte as a normal ASCII char
        //       , b => { self.write_byte(b); }
        //     }
        // }

    }

}

struct AnsiEscapeIter<'a> { curr_slice: &'a str
                          , in_escape: bool
                          }

impl<'a> AnsiEscapeIter<'a> {

    pub fn new(s: &'a str) -> Self {
        AnsiEscapeIter { curr_slice: s
                       , in_escape: false
                       }
    }
}

impl<'a> Iterator for AnsiEscapeIter<'a> {
    type Item = &'a str;

    fn next(&mut self) -> Option<Self::Item> {
        if self.curr_slice.len() == 0 {
            // if the remaining string is empty, we just return `None`
            None
        } else {
            // otherwise, find the next index to chunk on.
            let maybe_idx
                = if self.in_escape {
                     // if we're in an escape code, we split the chunk at the
                     // index of the next 'm' character, adding 1 so that the
                     // 'm' is in the escape code chunk.
                    self.curr_slice.find('m')
                        .map(|idx| idx + 1)
                } else {
                    // otherwise, split at the next ANSI escape sequence
                    self.curr_slice.find(ANSI_ESCAPE)
                };

            // if we found another index to chunk on, map over that index;
            // otherwise, we just yield the rest of the string
            maybe_idx.map_or(
                Some(self.curr_slice) // remainder (if no index to chunk on)
              , |idx| { // otherwise, chunk along that index...
                    let (chunk, next_slice) = self.curr_slice
                                                  .split_at(idx);
                    self.curr_slice = next_slice; // update current chunk
                    Some(chunk)                   // return the chunk
                })
        }

    }
}

impl Write for Terminal {

    fn write_str(&mut self, s: &str) -> Result {

        if s.contains(ANSI_ESCAPE) {
            // if the segment contains an ANSI escape, construct an iterator
            // over each chunk containing either an escape sequence or text
            for segment in AnsiEscapeIter::new(s) {
                if segment.starts_with(ANSI_ESCAPE) {
                    // if the current segment is an ANSI escape code,
                    // try to handle the escape and fail if it is malformed
                    self.handle_ansi_escape(segment)?
                } else {
                    // otherwise, just write each chunk in the string.
                    for byte in segment.as_bytes() {
                        self.write_byte(*byte);
                    }
                }
            }
        } else {
            // otherwise, if there are no ANSI escape codes,
            // we can just write each byte in the string.
            for byte in s.as_bytes() {
                self.write_byte(*byte);
            }
        }
        Ok(())
    }

}

/// Macro for printing to the standard output.
///
/// Equivalent to the `println!` macro except that a newline is not printed at
/// the end of the message.
///
/// # Panics
///
/// Panics if writing to the VGA console fails.
///
#[cfg(feature = "system_term")]
#[macro_export]
macro_rules! println {
    ($fmt:expr) => (print!(concat!($fmt, "\n")));
    ($fmt:expr, $($arg:tt)*) => (print!(concat!($fmt, "\n"), $($arg)*));
}

/// Macro for printing to the standard output.
///
/// Equivalent to the `println!` macro except that a newline is not printed at
/// the end of the message.
///
/// # Panics
///
/// Panics if writing to the VGA console fails.
#[cfg(feature = "system_term")]
#[macro_export]
macro_rules! print {
    ($($arg:tt)*) => ({
            use core::fmt::Write;
            $crate::CONSOLE.lock()
                           .write_fmt(format_args!($($arg)*))
                           .unwrap();
    });
}

#[cfg(feature = "log")]
#[macro_export]
macro_rules! log {
    (level: $level:expr, $fmt:expr, $($arg:tt)*) => ({
        println!( "[{} {}:{}] {}", $level, file!(), line!()
                , format_args!($fmt, $($arg)*));
    });
    (level: $level:expr, $fmt:expr) => ({
        println!("[{} {}:{}] {}", $level, file!(), line!(), $fmt);
    });
}
<commit_msg>[vga] note to self<commit_after>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! # SOS VGA
//! Code for interacting with the system's VGA buffer.
//!
//! # Features
//! This crate exposes the following feature flags:
//! + `system_term`: Use the VGA buffer as the system terminal. This provides
//!   implementations of the `print!` and `println!` macros.
//! + `log`: provides a `log!` macro for kernel log messages.
#![crate_name = "sos_vga"]
#![crate_type = "lib"]
#![feature( const_fn
          , slice_patterns
          , unique
          , question_mark )]
#![cfg_attr(feature = "system_term", feature(lang_items))]
#![cfg_attr(feature = "clippy", feature(plugin))]
#![cfg_attr(feature = "clippy", plugin(clippy))]
#![no_std]

#[cfg(feature = "system_term")]
extern crate spin;

use core::{mem, ptr};
use core::fmt::{Write, Result};
// use core::ptr::Unique;

#[cfg(feature = "system_term")]
use spin::Mutex;

/// The system's global VGA terminal
/// TODO: should this live in the kernel instead?
#[cfg(feature = "system_term")]
pub static CONSOLE: Mutex<Terminal>
    = Mutex::new(unsafe { Terminal::new(
         Palette::new(Color::LightGrey, Color::Black )
       , 0xB8000
    )});

// TODO: support varying VGA screen sizes?
pub const X_MAX: usize = 80;
pub const Y_MAX: usize = 25;

/// The type signature fot the actual VGA buffer
pub type Buffer = [[Char; X_MAX]; Y_MAX];

const ANSI_ESCAPE: &'static str = "\x1b";
const FG_MASK: u8 = 0b0000_1111;
const BG_MASK: u8 = 0b1111_0000;

/// VGA color codes
#[derive(Eq, PartialEq, Ord, PartialOrd, Copy, Clone)]
#[repr(u8)]
pub enum Color { Black = 0
               , Blue = 1
               , Green = 2
               , Cyan = 3
               , Red = 4
               , Magenta = 5
               , Brown = 6
               , LightGrey = 7
               , DarkGrey = 8
               , LightBlue = 9
               , LightGreen = 10
               , LightCyan = 11
               , LightRed = 12
               , LightMagenta = 13
               , Yellow = 14
               , White = 15
               }

#[derive(Eq, PartialEq, Ord, PartialOrd, Copy, Clone)]
pub struct Palette(u8);

impl Palette {
    /// Returns a `Palette` with the given foreground and background color.
    pub const fn new(fg: Color, bg: Color) -> Self {
        Palette( (bg as u8) << 4 | (fg as u8) )
    }

    /// Returns a new `Palette` with this palette's background color, and
    /// the specified foreground color.
    pub fn set_foreground(&self, fg: Color) -> Self {
        Palette( ( self.0 & BG_MASK) | (fg as u8 & FG_MASK) )
    }

    /// Returns a new `Palette` with this palette's foreground color, and
    /// the specified background color.
    pub fn set_background(&self, bg: Color) -> Self {
        Palette( ( (bg as u8) << 4 & BG_MASK) | (self.0 & FG_MASK) )
    }

    /// Returns this `Palette`'s foreground color.
    pub fn foreground(&self) -> Color {
        unsafe { mem::transmute(self.0 & FG_MASK) }
    }

    /// Returns this `Palette`'s background color.
    pub fn background(&self) -> Color {
        unsafe { mem::transmute((self.0 & BG_MASK) >> 4) }
    }

}
/// A colored VGA character.
#[derive(Copy, Clone)]
#[repr(C)]
pub struct Char { pub ascii: u8
                , pub colors: Palette
                }

impl Char {
    pub const fn empty() -> Self {
        Char { ascii: 0, colors: Palette(0) }
    }
}

pub struct Terminal { buffer: ptr::Unique<Buffer>
                    , x: usize
                    , y: usize
                    , colors: Palette
                    }
impl Terminal {

    /// Constructs a new `Terminal` for abuffer starting at the given address.
    ///
    /// # Arguments:
    /// + `colors`: the default color palette for the terminal
    /// + `buffer_start`: the address of the to the memory location where
    ///                   the terminal's VGA buffer begins
    ///
    /// # Unsafe due to:
    /// + Casting a raw address to an array
    pub const unsafe fn new( colors: Palette
                           , buffer_start: usize)
                           -> Terminal {
        Terminal { x: 0, y: 0
                 , colors: colors
                 , buffer: ptr::Unique::new(buffer_start as *mut _)
                 }
    }

    #[inline]
    fn buffer(&mut self) -> &mut  Buffer {
        unsafe { self.buffer.get_mut() }
    }

    /// Set the color palette used for writing subsequent characters.
    pub fn set_colors(&mut self, bg: Color, fg: Color) -> &mut Self {
        self.colors = Palette::new(bg,fg);
        self
    }

    /// Scrolls the terminal one row.
    fn scroll(&mut self) {
        // // construct an iterator over the whole buffer
        // let mut rows = self.buffer()
        //                    .iter_mut();
        //
        // // the current row in the buffer
        // let mut current = rows.next()
        //                    .unwrap();
        //
        // while let Some(next) = rows.next() {
        //     // while there are rows remaining in the iterator, swap the
        //     // next row with the current one (moving it back by one)
        //     mem::swap(current, next);
        //     // and advance our pointer to the current row.
        //     current = next;
        // }
        let buffer: &mut _ = self.buffer();
        for row in 1..Y_MAX {
            for col in 0 .. X_MAX {
                unsafe {
                    let last = ptr::read_volatile(&buffer[row][col]);
                    ptr::write_volatile( &mut buffer[row - 1][col], last);
                }
            }
        }
        // empty the last line in the buffer
        unsafe {
            for col in 0..X_MAX {
                ptr::write_volatile(&mut buffer[Y_MAX - 1][col], Char::empty());
            }
        }
    }

    /// Clear the terminal
    pub fn clear(&mut self) -> &mut Self {
        // to clear the terminal, we just zero out the whole buffer.
        {
            let buffer: &mut _ =self.buffer();
            for row in 0..Y_MAX {
                for col in 0 .. X_MAX {
                    unsafe {
                        ptr::write_volatile(
                            &mut buffer[row][col], Char::empty()
                        );
                    }
                }
            }
        }

        self.x = 0;
        self.y = 0;
        self
    }

    /// Write the given byte to the terminal, and advance the cursor position.
    pub fn write_byte(&mut self, byte: u8) -> &mut Self {
        if byte == b'\n' {
            // if the byte is a newline, we just advance to the next line
            // and reset the column position.
            self.x = 0;
            self.y += 1;
        } else {
            // otherwise, it's a regular character, so we just set the
            // byte at the current position in the buffer to that
            // character (with the current color palette)
            //
            let y = self.y;
            let x = self.x;
            unsafe {
                ptr::write_volatile( &mut (self.buffer()[y][x])
                                   , Char { ascii: byte, colors: self.colors }
                                   );
            }
            // and advance our column position by one
            self.x += 1;

            if self.x >= X_MAX {
                // if we've reached the end of the line, advance to the next
                self.x = 0;
                self.y += 1;
            }
        }
        if self.y >= Y_MAX {
            // if we've reached the bottom of the terminal, scroll.
            self.scroll();
            self.y -= 1;
        }
        self
    }

    fn handle_ansi_escape(&self, escape_code: &str) -> Result {
        match escape_code.as_bytes() {
            // `\x1b[3Nm` sets the foreground color to N.
            &[0x1b, b'[', b'3', n, b'm'] => {
                unsafe { self.colors
                             .set_foreground(mem::transmute(n - 48)); }
                Ok(())
            }
            // `\x1b[4Nm` sets the background color to N
          , &[0x1b, b'[', b'4', n, b'm'] => {
                unsafe { self.colors
                             .set_background(mem::transmute(n - 48)); }
                Ok(())
            }
          , _ => unimplemented!()
        }
        // let escape_seq: &str = bytes.take_while(|b| b != b'm')
        //                       .collect::<&str>();
        // match escape_seq {
        //     [b'3', n] => unsafe {
        //         self.colors.set_foreground(mem::transmute(n - 48))
        //     }
        // }
        // while let Some(byte) = bytes.next() {
        //     match *byte {
        //         // we've recieved an ANSI escape sequence.
        //         // this basically enters a mediocre FSM for matching ANSI
        //         // control codes.
        //         0x1b => match *next_ansi_byte!(bytes) {
        //             // handle multi-char ANSI escapes
        //             b'[' => match *next_ansi_byte!(bytes) {
        //                 // foreground color code
        //                 fg @ 30 ... 37 => {
        //                     if !(*next_ansi_byte!(bytes) == b'm') {
        //                         unsafe {
        //                             let color: vga::Color
        //                                 = mem::transmute(fg - 30);
        //                             self.colors
        //                                 .set_foreground(color);
        //                         }
        //
        //                     }
        //                 }
        //                 // background color code
        //               , 40 ... 47 => {
        //
        //                 }
        //               , _ => unimplemented!()
        //             }
        //           , _    => unimplemented!()
        //         }
        //         // otherwise, treat the byte as a normal ASCII char
        //       , b => { self.write_byte(b); }
        //     }
        // }

    }

}

struct AnsiEscapeIter<'a> { curr_slice: &'a str
                          , in_escape: bool
                          }

impl<'a> AnsiEscapeIter<'a> {

    pub fn new(s: &'a str) -> Self {
        AnsiEscapeIter { curr_slice: s
                       , in_escape: false
                       }
    }
}

impl<'a> Iterator for AnsiEscapeIter<'a> {
    type Item = &'a str;

    fn next(&mut self) -> Option<Self::Item> {
        if self.curr_slice.len() == 0 {
            // if the remaining string is empty, we just return `None`
            None
        } else {
            // otherwise, find the next index to chunk on.
            let maybe_idx
                = if self.in_escape {
                     // if we're in an escape code, we split the chunk at the
                     // index of the next 'm' character, adding 1 so that the
                     // 'm' is in the escape code chunk.
                    self.curr_slice.find('m')
                        .map(|idx| idx + 1)
                } else {
                    // otherwise, split at the next ANSI escape sequence
                    self.curr_slice.find(ANSI_ESCAPE)
                };

            // if we found another index to chunk on, map over that index;
            // otherwise, we just yield the rest of the string
            maybe_idx.map_or(
                Some(self.curr_slice) // remainder (if no index to chunk on)
              , |idx| { // otherwise, chunk along that index...
                    let (chunk, next_slice) = self.curr_slice
                                                  .split_at(idx);
                    self.curr_slice = next_slice; // update current chunk
                    Some(chunk)                   // return the chunk
                })
        }

    }
}

impl Write for Terminal {

    fn write_str(&mut self, s: &str) -> Result {

        if s.contains(ANSI_ESCAPE) {
            // if the segment contains an ANSI escape, construct an iterator
            // over each chunk containing either an escape sequence or text
            for segment in AnsiEscapeIter::new(s) {
                if segment.starts_with(ANSI_ESCAPE) {
                    // if the current segment is an ANSI escape code,
                    // try to handle the escape and fail if it is malformed
                    self.handle_ansi_escape(segment)?
                } else {
                    // otherwise, just write each chunk in the string.
                    for byte in segment.as_bytes() {
                        self.write_byte(*byte);
                    }
                }
            }
        } else {
            // otherwise, if there are no ANSI escape codes,
            // we can just write each byte in the string.
            for byte in s.as_bytes() {
                self.write_byte(*byte);
            }
        }
        Ok(())
    }

}

/// Macro for printing to the standard output.
///
/// Equivalent to the `println!` macro except that a newline is not printed at
/// the end of the message.
///
/// # Panics
///
/// Panics if writing to the VGA console fails.
///
#[cfg(feature = "system_term")]
#[macro_export]
macro_rules! println {
    ($fmt:expr) => (print!(concat!($fmt, "\n")));
    ($fmt:expr, $($arg:tt)*) => (print!(concat!($fmt, "\n"), $($arg)*));
}

/// Macro for printing to the standard output.
///
/// Equivalent to the `println!` macro except that a newline is not printed at
/// the end of the message.
///
/// # Panics
///
/// Panics if writing to the VGA console fails.
#[cfg(feature = "system_term")]
#[macro_export]
macro_rules! print {
    ($($arg:tt)*) => ({
            use core::fmt::Write;
            $crate::CONSOLE.lock()
                           .write_fmt(format_args!($($arg)*))
                           .unwrap();
    });
}

#[cfg(feature = "log")]
#[macro_export]
macro_rules! log {
    (level: $level:expr, $fmt:expr, $($arg:tt)*) => ({
        println!( "[{} {}:{}] {}", $level, file!(), line!()
                , format_args!($fmt, $($arg)*));
    });
    (level: $level:expr, $fmt:expr) => ({
        println!("[{} {}:{}] {}", $level, file!(), line!(), $fmt);
    });
}
<|endoftext|>


<commit_before>// Copyright 2016 Alexander Reece
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

mod class_mod;
mod frame_payload_enum;
mod header_enum;
mod method_enum;

use std::iter::ExactSizeIterator;
use std::io;

use WriteRust;
use common::Spec;

use self::class_mod::ClassModuleWriter;
use self::frame_payload_enum::FramePayloadEnumWriter;
use self::header_enum::HeaderEnumWriter;
use self::method_enum::MethodEnumWriter;

pub struct SpecModuleWriter<'a> {
    spec: &'a Spec,
}

impl<'a> WriteRust for SpecModuleWriter<'a> {
    fn write_rust_to<W>(&self, writer: &mut W) -> io::Result<()>
        where W: io::Write
    {
        if self.spec.classes().len() == 0 {
            return Ok(());
        }

        try!(writeln!(writer, "#![allow(too_many_arguments)]\n"));

        try!(self.write_class_constants(writer));
        try!(self.write_method_constants(writer));

        let header_enum = HeaderEnumWriter::new(&self.spec);
        try!(header_enum.write_rust_to(writer));
        try!(self.write_frame_struct(writer));

        let frame_payload_enum = FramePayloadEnumWriter::new(&self.spec);
        try!(frame_payload_enum.write_rust_to(writer));

        try!(writeln!(writer, "\n// Class Modules"));
        for class in self.spec.classes() {
            let module_writer = ClassModuleWriter::new(class);
            try!(module_writer.write_rust_to(writer));
        }

        // aliases
        try!(writeln!(writer, "\n// Class methods"));
        for class in self.spec.classes() {
            try!(writeln!(writer, "pub use self::{}::ClassMethod as {}Method;", class.snake_case(), class.pascal_case()));
        }

        try!(writeln!(writer, "\n// Class headers"));
        for class in self.spec.classes() {
            if class.fields().is_empty() {
                continue;
            }
            try!(writeln!(writer, "pub use self::{}::Header as {}Header;", class.snake_case(), class.pascal_case()));
        }

        let method_enum = MethodEnumWriter::new(&self.spec);
        try!(method_enum.write_rust_to(writer));

        Ok(())
    }
}

impl<'a> SpecModuleWriter<'a> {
    pub fn new(spec: &'a Spec) -> Self {
        SpecModuleWriter {
            spec: spec,
        }
    }

    pub fn mod_name(&self) -> &str {
        self.spec.mod_name()
    }

    pub fn write_class_constants<W>(&self, writer: &mut W) -> io::Result<()>
        where W: io::Write
    {
        try!(writeln!(writer, "// Class Constants"));
        try!(writeln!(writer, "// Generated by primalgen::codegen::spec_module::SpecModuleWriter"));
        for (_, class_names) in self.spec.class_indexes() {
            for class_name in class_names {
                let class = self.spec.class(class_name).unwrap();
                let constant_class = class.constant_case();
                try!(writeln!(writer, "pub const CLASS_{}: u16 = {};", constant_class, class.index()));
            }
        }

        Ok(())
    }

    /// Write the correct frame strucutre
    /// 0.8.x had a "cycle" in each frame, 0.9.x doesn't
    pub fn write_frame_struct<W>(&self, writer: &mut W) -> io::Result<()>
        where W: io::Write
    {
        let version = self.spec.version();

        let (cycle_prop, cycle_getter, cycle_arg, cycle_create) = if version.minor() == 8 {
            (
                "cycle: u8,\n",
                "\n\
                    pub fn cycle(&self) -> u8 {{\n\
                        self.cycle\n\
                    }} // fn cycle()\n\
                 \n",
                "cycle: u8,\n",
                "cycle: cycle,\n",
            )
        } else {
            ("", "", "", "")
        };

        // struct Frame
        try!(writeln!(writer, "\n\
                // Generated by primalgen::codegen::spec_module::SpecModuleWriter
                #[derive(Debug)]
                pub struct Frame<'a> {{\n\
                    channel: u16,\n\
                    {}\
                    payload: FramePayload<'a>,\n\
                }} // struct Frame\n\
                \n\
                impl<'a> Frame<'a> {{\n\
                    pub fn new<P>(channel: u16, {}payload: P) -> Self\n\
                        where P: Into<FramePayload<'a>>\n\
                    {{\n\
                        Frame {{\n\
                            channel: channel,\n\
                            {}\
                            payload: payload.into(),\n\
                        }} // Frame\n\
                    }} // fn new()\n\
                    \n\
                    pub fn channel(&self) -> u16 {{\n\
                        self.channel\n\
                    }} // fn channel()\n\
                    \n\
                    pub fn payload(&self) -> &FramePayload<'a> {{\n\
                        &self.payload\n\
                    }} // fn payload()\n\
                    {}\n\
                 }} // impl Frame<'a>
                    ",
                      cycle_prop, // struct Frame<'a>...
                      cycle_arg, // fn new()...
                      cycle_create, // Frame { cycle: ... }
                      cycle_getter, // fn cycle() -> u8
        ));

        Ok(())
    }

    pub fn write_method_constants<W>(&self, writer: &mut W) -> io::Result<()>
        where W: io::Write
    {
        try!(writeln!(writer, "\n// Class Methods"));
        try!(write!(writer, "// Generated by codegen::spec_module::SpecModuleWriter"));
        for class in self.spec.classes() {
            try!(writeln!(writer, ""));
            let constant_class = class.constant_case();
            for (_, method_names) in class.method_indexes() {
                for method_name in method_names {
                    let method = class.method(method_name).unwrap();
                    try!(writeln!(
                        writer,
                        "pub const METHOD_{}_{}: u16 = {};",
                        class.constant_case(),
                        method.constant_case(),
                        method.index()
                    ));
                }
            }
        }

        Ok(())
    }
}<commit_msg>[primalgen] Remove unused variable from SpecModuleWriter<commit_after>// Copyright 2016 Alexander Reece
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

mod class_mod;
mod frame_payload_enum;
mod header_enum;
mod method_enum;

use std::iter::ExactSizeIterator;
use std::io;

use WriteRust;
use common::Spec;

use self::class_mod::ClassModuleWriter;
use self::frame_payload_enum::FramePayloadEnumWriter;
use self::header_enum::HeaderEnumWriter;
use self::method_enum::MethodEnumWriter;

pub struct SpecModuleWriter<'a> {
    spec: &'a Spec,
}

impl<'a> WriteRust for SpecModuleWriter<'a> {
    fn write_rust_to<W>(&self, writer: &mut W) -> io::Result<()>
        where W: io::Write
    {
        if self.spec.classes().len() == 0 {
            return Ok(());
        }

        try!(writeln!(writer, "#![allow(too_many_arguments)]\n"));

        try!(self.write_class_constants(writer));
        try!(self.write_method_constants(writer));

        let header_enum = HeaderEnumWriter::new(&self.spec);
        try!(header_enum.write_rust_to(writer));
        try!(self.write_frame_struct(writer));

        let frame_payload_enum = FramePayloadEnumWriter::new(&self.spec);
        try!(frame_payload_enum.write_rust_to(writer));

        try!(writeln!(writer, "\n// Class Modules"));
        for class in self.spec.classes() {
            let module_writer = ClassModuleWriter::new(class);
            try!(module_writer.write_rust_to(writer));
        }

        // aliases
        try!(writeln!(writer, "\n// Class methods"));
        for class in self.spec.classes() {
            try!(writeln!(writer, "pub use self::{}::ClassMethod as {}Method;", class.snake_case(), class.pascal_case()));
        }

        try!(writeln!(writer, "\n// Class headers"));
        for class in self.spec.classes() {
            if class.fields().is_empty() {
                continue;
            }
            try!(writeln!(writer, "pub use self::{}::Header as {}Header;", class.snake_case(), class.pascal_case()));
        }

        let method_enum = MethodEnumWriter::new(&self.spec);
        try!(method_enum.write_rust_to(writer));

        Ok(())
    }
}

impl<'a> SpecModuleWriter<'a> {
    pub fn new(spec: &'a Spec) -> Self {
        SpecModuleWriter {
            spec: spec,
        }
    }

    pub fn mod_name(&self) -> &str {
        self.spec.mod_name()
    }

    pub fn write_class_constants<W>(&self, writer: &mut W) -> io::Result<()>
        where W: io::Write
    {
        try!(writeln!(writer, "// Class Constants"));
        try!(writeln!(writer, "// Generated by primalgen::codegen::spec_module::SpecModuleWriter"));
        for (_, class_names) in self.spec.class_indexes() {
            for class_name in class_names {
                let class = self.spec.class(class_name).unwrap();
                let constant_class = class.constant_case();
                try!(writeln!(writer, "pub const CLASS_{}: u16 = {};", constant_class, class.index()));
            }
        }

        Ok(())
    }

    /// Write the correct frame strucutre
    /// 0.8.x had a "cycle" in each frame, 0.9.x doesn't
    pub fn write_frame_struct<W>(&self, writer: &mut W) -> io::Result<()>
        where W: io::Write
    {
        let version = self.spec.version();

        let (cycle_prop, cycle_getter, cycle_arg, cycle_create) = if version.minor() == 8 {
            (
                "cycle: u8,\n",
                "\n\
                    pub fn cycle(&self) -> u8 {{\n\
                        self.cycle\n\
                    }} // fn cycle()\n\
                 \n",
                "cycle: u8,\n",
                "cycle: cycle,\n",
            )
        } else {
            ("", "", "", "")
        };

        // struct Frame
        try!(writeln!(writer, "\n\
                // Generated by primalgen::codegen::spec_module::SpecModuleWriter
                #[derive(Debug)]
                pub struct Frame<'a> {{\n\
                    channel: u16,\n\
                    {}\
                    payload: FramePayload<'a>,\n\
                }} // struct Frame\n\
                \n\
                impl<'a> Frame<'a> {{\n\
                    pub fn new<P>(channel: u16, {}payload: P) -> Self\n\
                        where P: Into<FramePayload<'a>>\n\
                    {{\n\
                        Frame {{\n\
                            channel: channel,\n\
                            {}\
                            payload: payload.into(),\n\
                        }} // Frame\n\
                    }} // fn new()\n\
                    \n\
                    pub fn channel(&self) -> u16 {{\n\
                        self.channel\n\
                    }} // fn channel()\n\
                    \n\
                    pub fn payload(&self) -> &FramePayload<'a> {{\n\
                        &self.payload\n\
                    }} // fn payload()\n\
                    {}\n\
                 }} // impl Frame<'a>
                    ",
                      cycle_prop, // struct Frame<'a>...
                      cycle_arg, // fn new()...
                      cycle_create, // Frame { cycle: ... }
                      cycle_getter, // fn cycle() -> u8
        ));

        Ok(())
    }

    pub fn write_method_constants<W>(&self, writer: &mut W) -> io::Result<()>
        where W: io::Write
    {
        try!(writeln!(writer, "\n// Class Methods"));
        try!(write!(writer, "// Generated by codegen::spec_module::SpecModuleWriter"));
        for class in self.spec.classes() {
            try!(writeln!(writer, ""));
            for (_, method_names) in class.method_indexes() {
                for method_name in method_names {
                    let method = class.method(method_name).unwrap();
                    try!(writeln!(
                        writer,
                        "pub const METHOD_{}_{}: u16 = {};",
                        class.constant_case(),
                        method.constant_case(),
                        method.index()
                    ));
                }
            }
        }

        Ok(())
    }
}<|endoftext|>


<commit_before>use core_simd::*;

use std::f64::consts::PI;
const SOLAR_MASS: f64 = 4.0 * PI * PI;
const DAYS_PER_YEAR: f64 = 365.24;

#[derive(Debug, Clone, Copy)]
pub struct Body {
    pub x: f64x4,
    pub v: f64x4,
    pub mass: f64,
}

// translation up to here
const N_BODIES: usize = 5;
#[allow(clippy::unreadable_literal)]
const BODIES: [Body; N_BODIES] = [
    // sun:
    Body {
        x: f64x4::from_array([0., 0., 0., 0.]),
        v: f64x4::from_array([0., 0., 0., 0.]),
        mass: SOLAR_MASS,
    },
    // jupiter:
    Body {
        x: f64x4::from_array([
            4.84143144246472090e+00,
            -1.16032004402742839e+00,
            -1.03622044471123109e-01,
            0.,
        ]),
        v: f64x4::from_array([
            1.66007664274403694e-03 * DAYS_PER_YEAR,
            7.69901118419740425e-03 * DAYS_PER_YEAR,
            -6.90460016972063023e-05 * DAYS_PER_YEAR,
            0.,
        ]),
        mass: 9.54791938424326609e-04 * SOLAR_MASS,
    },
    // saturn:
    Body {
        x: f64x4::from_array([
            8.34336671824457987e+00,
            4.12479856412430479e+00,
            -4.03523417114321381e-01,
            0.,
        ]),
        v: f64x4::from_array([
            -2.76742510726862411e-03 * DAYS_PER_YEAR,
            4.99852801234917238e-03 * DAYS_PER_YEAR,
            2.30417297573763929e-05 * DAYS_PER_YEAR,
            0.,
        ]),
        mass: 2.85885980666130812e-04 * SOLAR_MASS,
    },
    // uranus:
    Body {
        x: f64x4::from_array([
            1.28943695621391310e+01,
            -1.51111514016986312e+01,
            -2.23307578892655734e-01,
            0.,
        ]),
        v: f64x4::from_array([
            2.96460137564761618e-03 * DAYS_PER_YEAR,
            2.37847173959480950e-03 * DAYS_PER_YEAR,
            -2.96589568540237556e-05 * DAYS_PER_YEAR,
            0.,
        ]),
        mass: 4.36624404335156298e-05 * SOLAR_MASS,
    },
    // neptune:
    Body {
        x: f64x4::from_array([
            1.53796971148509165e+01,
            -2.59193146099879641e+01,
            1.79258772950371181e-01,
            0.,
        ]),
        v: f64x4::from_array([
            2.68067772490389322e-03 * DAYS_PER_YEAR,
            1.62824170038242295e-03 * DAYS_PER_YEAR,
            -9.51592254519715870e-05 * DAYS_PER_YEAR,
            0.,
        ]),
        mass: 5.15138902046611451e-05 * SOLAR_MASS,
    },
];

pub fn offset_momentum(bodies: &mut [Body; N_BODIES]) {
    let (sun, rest) = bodies.split_at_mut(1);
    let sun = &mut sun[0];
    for body in rest {
        let m_ratio = body.mass / SOLAR_MASS;
        sun.v -= body.v * m_ratio;
    }
}

pub fn energy(bodies: &[Body; N_BODIES]) -> f64 {
    let mut e = 0.;
    for i in 0..N_BODIES {
        let bi = &bodies[i];
        e += bi.mass * (bi.v * bi.v).horizontal_sum() * 0.5;
        for bj in bodies.iter().take(N_BODIES).skip(i + 1) {
            let dx = bi.x - bj.x;
            e -= bi.mass * bj.mass / (dx * dx).horizontal_sum().sqrt()
        }
    }
    e
}

pub fn advance(bodies: &mut [Body; N_BODIES], dt: f64) {
    const N: usize = N_BODIES * (N_BODIES - 1) / 2;

    // compute distance between bodies:
    let mut r = [f64x4::splat(0.); N];
    {
        let mut i = 0;
        for j in 0..N_BODIES {
            for k in j + 1..N_BODIES {
                r[i] = bodies[j].x - bodies[k].x;
                i += 1;
            }
        }
    }

    let mut mag = [0.0; N];
    let mut i = 0;
    while i < N {
        let d2s = f64x2::from_array([(r[i] * r[i]).horizontal_sum(), (r[i + 1] * r[i + 1]).horizontal_sum()]);
        let dmags = f64x2::splat(dt) / (d2s * d2s.sqrt());
        dmags.write_to_slice_unaligned(&mut mag[i..]);
        i += 2;
    }

    i = 0;
    for j in 0..N_BODIES {
        for k in j + 1..N_BODIES {
            let f = r[i] * mag[i];
            bodies[j].v -= f * bodies[k].mass;
            bodies[k].v += f * bodies[j].mass;
            i += 1
        }
    }
    for body in bodies {
        body.x += dt * body.v
    }
}

pub fn run_k<K>(n: usize, k: K) -> (f64, f64)
where
    K: Fn(&mut [Body; N_BODIES], f64),
{
    let mut bodies = BODIES;
    offset_momentum(&mut bodies);
    let energy_before = energy(&bodies);
    for _ in 0..n {
        k(&mut bodies, 0.01);
    }
    let energy_after = energy(&bodies);

    (energy_before, energy_after)
}

pub fn run(n: usize) -> (f64, f64) {
    run_k(n, advance)
}

const OUTPUT: Vec<f64> = vec![-0.169075164, -0.169087605];
#[cfg(test)]
mod tests {
    #[test]
    fn test() {
        let mut out: Vec<u8> = Vec::new();
        run(&mut out, 1000, 0);
        for &(size, a_e, b_e) in crate::RESULTS {
            let (a, b) = super::run(size);
            assert_eq!(format!("{:.9}", a), a_e);
            assert_eq!(format!("{:.9}", b), b_e);
        }
    }
}
fn main() {
    //let n: usize = std::env::args()
    //.nth(1)
    //.expect("need one arg")
    //.parse()
    //.expect("argument should be a usize");
    //run(&mut std::io::stdout(), n, alg);
    println!("{:?}", run_k::<10>(10, 10));
}
<commit_msg>don't need clippy<commit_after>use core_simd::*;

use std::f64::consts::PI;
const SOLAR_MASS: f64 = 4.0 * PI * PI;
const DAYS_PER_YEAR: f64 = 365.24;

#[derive(Debug, Clone, Copy)]
pub struct Body {
    pub x: f64x4,
    pub v: f64x4,
    pub mass: f64,
}

// translation up to here
const N_BODIES: usize = 5;
const BODIES: [Body; N_BODIES] = [
    // sun:
    Body {
        x: f64x4::from_array([0., 0., 0., 0.]),
        v: f64x4::from_array([0., 0., 0., 0.]),
        mass: SOLAR_MASS,
    },
    // jupiter:
    Body {
        x: f64x4::from_array([
            4.84143144246472090e+00,
            -1.16032004402742839e+00,
            -1.03622044471123109e-01,
            0.,
        ]),
        v: f64x4::from_array([
            1.66007664274403694e-03 * DAYS_PER_YEAR,
            7.69901118419740425e-03 * DAYS_PER_YEAR,
            -6.90460016972063023e-05 * DAYS_PER_YEAR,
            0.,
        ]),
        mass: 9.54791938424326609e-04 * SOLAR_MASS,
    },
    // saturn:
    Body {
        x: f64x4::from_array([
            8.34336671824457987e+00,
            4.12479856412430479e+00,
            -4.03523417114321381e-01,
            0.,
        ]),
        v: f64x4::from_array([
            -2.76742510726862411e-03 * DAYS_PER_YEAR,
            4.99852801234917238e-03 * DAYS_PER_YEAR,
            2.30417297573763929e-05 * DAYS_PER_YEAR,
            0.,
        ]),
        mass: 2.85885980666130812e-04 * SOLAR_MASS,
    },
    // uranus:
    Body {
        x: f64x4::from_array([
            1.28943695621391310e+01,
            -1.51111514016986312e+01,
            -2.23307578892655734e-01,
            0.,
        ]),
        v: f64x4::from_array([
            2.96460137564761618e-03 * DAYS_PER_YEAR,
            2.37847173959480950e-03 * DAYS_PER_YEAR,
            -2.96589568540237556e-05 * DAYS_PER_YEAR,
            0.,
        ]),
        mass: 4.36624404335156298e-05 * SOLAR_MASS,
    },
    // neptune:
    Body {
        x: f64x4::from_array([
            1.53796971148509165e+01,
            -2.59193146099879641e+01,
            1.79258772950371181e-01,
            0.,
        ]),
        v: f64x4::from_array([
            2.68067772490389322e-03 * DAYS_PER_YEAR,
            1.62824170038242295e-03 * DAYS_PER_YEAR,
            -9.51592254519715870e-05 * DAYS_PER_YEAR,
            0.,
        ]),
        mass: 5.15138902046611451e-05 * SOLAR_MASS,
    },
];

pub fn offset_momentum(bodies: &mut [Body; N_BODIES]) {
    let (sun, rest) = bodies.split_at_mut(1);
    let sun = &mut sun[0];
    for body in rest {
        let m_ratio = body.mass / SOLAR_MASS;
        sun.v -= body.v * m_ratio;
    }
}

pub fn energy(bodies: &[Body; N_BODIES]) -> f64 {
    let mut e = 0.;
    for i in 0..N_BODIES {
        let bi = &bodies[i];
        e += bi.mass * (bi.v * bi.v).horizontal_sum() * 0.5;
        for bj in bodies.iter().take(N_BODIES).skip(i + 1) {
            let dx = bi.x - bj.x;
            e -= bi.mass * bj.mass / (dx * dx).horizontal_sum().sqrt()
        }
    }
    e
}

pub fn advance(bodies: &mut [Body; N_BODIES], dt: f64) {
    const N: usize = N_BODIES * (N_BODIES - 1) / 2;

    // compute distance between bodies:
    let mut r = [f64x4::splat(0.); N];
    {
        let mut i = 0;
        for j in 0..N_BODIES {
            for k in j + 1..N_BODIES {
                r[i] = bodies[j].x - bodies[k].x;
                i += 1;
            }
        }
    }

    let mut mag = [0.0; N];
    let mut i = 0;
    while i < N {
        let d2s = f64x2::from_array([(r[i] * r[i]).horizontal_sum(), (r[i + 1] * r[i + 1]).horizontal_sum()]);
        let dmags = f64x2::splat(dt) / (d2s * d2s.sqrt());
        dmags.write_to_slice_unaligned(&mut mag[i..]);
        i += 2;
    }

    i = 0;
    for j in 0..N_BODIES {
        for k in j + 1..N_BODIES {
            let f = r[i] * mag[i];
            bodies[j].v -= f * bodies[k].mass;
            bodies[k].v += f * bodies[j].mass;
            i += 1
        }
    }
    for body in bodies {
        body.x += dt * body.v
    }
}

pub fn run_k<K>(n: usize, k: K) -> (f64, f64)
where
    K: Fn(&mut [Body; N_BODIES], f64),
{
    let mut bodies = BODIES;
    offset_momentum(&mut bodies);
    let energy_before = energy(&bodies);
    for _ in 0..n {
        k(&mut bodies, 0.01);
    }
    let energy_after = energy(&bodies);

    (energy_before, energy_after)
}

pub fn run(n: usize) -> (f64, f64) {
    run_k(n, advance)
}

const OUTPUT: Vec<f64> = vec![-0.169075164, -0.169087605];
#[cfg(test)]
mod tests {
    #[test]
    fn test() {
        let mut out: Vec<u8> = Vec::new();
        run(&mut out, 1000, 0);
        for &(size, a_e, b_e) in crate::RESULTS {
            let (a, b) = super::run(size);
            assert_eq!(format!("{:.9}", a), a_e);
            assert_eq!(format!("{:.9}", b), b_e);
        }
    }
}
fn main() {
    //let n: usize = std::env::args()
    //.nth(1)
    //.expect("need one arg")
    //.parse()
    //.expect("argument should be a usize");
    //run(&mut std::io::stdout(), n, alg);
    println!("{:?}", run_k::<10>(10, 10));
}
<|endoftext|>


<commit_before>//! Inference of closure parameter types based on the closure's expected type.

use chalk_ir::{fold::Shift, AliasTy, FnSubst, WhereClause};
use hir_def::HasModule;
use smallvec::SmallVec;

use crate::{
    to_chalk_trait_id, utils, ChalkTraitId, DynTy, FnPointer, FnSig, Interner, Substitution, Ty,
    TyKind,
};

use super::{Expectation, InferenceContext};

impl InferenceContext<'_> {
    pub(super) fn deduce_closure_type_from_expectations(
        &mut self,
        closure_ty: &Ty,
        sig_ty: &Ty,
        expectation: &Expectation,
    ) {
        let expected_ty = match expectation.to_option(&mut self.table) {
            Some(ty) => ty,
            None => return,
        };

        // Deduction from where-clauses in scope, as well as fn-pointer coercion are handled here.
        self.coerce(closure_ty, &expected_ty);

        // Deduction based on the expected `dyn Fn` is done separately.
        if let TyKind::Dyn(dyn_ty) = expected_ty.kind(&Interner) {
            if let Some(sig) = self.deduce_sig_from_dyn_ty(dyn_ty) {
                let expected_sig_ty = TyKind::Function(sig).intern(&Interner);

                self.unify(sig_ty, &expected_sig_ty);
            }
        }
    }

    fn deduce_sig_from_dyn_ty(&self, dyn_ty: &DynTy) -> Option<FnPointer> {
        // Search for predicates like `$self: FnX<Args>` and `<$self as FnOnce<...>>::Output == Ret`

        let fn_traits: SmallVec<[ChalkTraitId; 3]> =
            utils::fn_traits(self.db.upcast(), self.owner.module(self.db.upcast()).krate())
                .map(|tid| to_chalk_trait_id(tid))
                .collect();

        for bound in dyn_ty.bounds.map_ref(|b| b.iter(&Interner)) {
            let bound = bound.map(|b| b.clone()).fuse_binders(&Interner);
            match bound.skip_binders() {
                WhereClause::AliasEq(eq) => match &eq.alias {
                    AliasTy::Projection(projection) => {
                        let assoc_data = self.db.associated_ty_data(projection.associated_ty_id);
                        if !fn_traits.contains(&assoc_data.trait_id) {
                            return None;
                        }

                        // Skip `Self`, get the type argument.
                        let arg = projection.substitution.as_slice(&Interner).get(1)?;
                        match arg.ty(&Interner)?.kind(&Interner) {
                            TyKind::Tuple(_, subst) => {
                                let generic_args = subst.as_slice(&Interner);
                                let mut sig_tys = Vec::new();
                                for arg in generic_args {
                                    sig_tys.push(arg.ty(&Interner)?.clone());
                                }
                                sig_tys.push(eq.ty.clone());

                                cov_mark::hit!(dyn_fn_param_informs_call_site_closure_signature);
                                return Some(FnPointer {
                                    num_binders: 0,
                                    sig: FnSig {
                                        abi: (),
                                        safety: chalk_ir::Safety::Safe,
                                        variadic: false,
                                    },
                                    substitution: FnSubst(
                                        Substitution::from_iter(&Interner, sig_tys.clone())
                                            .shifted_in(&Interner),
                                    ),
                                });
                            }
                            _ => {}
                        }
                    }
                    AliasTy::Opaque(_) => {}
                },
                _ => {}
            }
        }

        None
    }
}
<commit_msg>update comment<commit_after>//! Inference of closure parameter types based on the closure's expected type.

use chalk_ir::{fold::Shift, AliasTy, FnSubst, WhereClause};
use hir_def::HasModule;
use smallvec::SmallVec;

use crate::{
    to_chalk_trait_id, utils, ChalkTraitId, DynTy, FnPointer, FnSig, Interner, Substitution, Ty,
    TyKind,
};

use super::{Expectation, InferenceContext};

impl InferenceContext<'_> {
    pub(super) fn deduce_closure_type_from_expectations(
        &mut self,
        closure_ty: &Ty,
        sig_ty: &Ty,
        expectation: &Expectation,
    ) {
        let expected_ty = match expectation.to_option(&mut self.table) {
            Some(ty) => ty,
            None => return,
        };

        // Deduction from where-clauses in scope, as well as fn-pointer coercion are handled here.
        self.coerce(closure_ty, &expected_ty);

        // Deduction based on the expected `dyn Fn` is done separately.
        if let TyKind::Dyn(dyn_ty) = expected_ty.kind(&Interner) {
            if let Some(sig) = self.deduce_sig_from_dyn_ty(dyn_ty) {
                let expected_sig_ty = TyKind::Function(sig).intern(&Interner);

                self.unify(sig_ty, &expected_sig_ty);
            }
        }
    }

    fn deduce_sig_from_dyn_ty(&self, dyn_ty: &DynTy) -> Option<FnPointer> {
        // Search for a predicate like `<$self as FnX<Args>>::Output == Ret`

        let fn_traits: SmallVec<[ChalkTraitId; 3]> =
            utils::fn_traits(self.db.upcast(), self.owner.module(self.db.upcast()).krate())
                .map(|tid| to_chalk_trait_id(tid))
                .collect();

        for bound in dyn_ty.bounds.map_ref(|b| b.iter(&Interner)) {
            let bound = bound.map(|b| b.clone()).fuse_binders(&Interner);
            match bound.skip_binders() {
                WhereClause::AliasEq(eq) => match &eq.alias {
                    AliasTy::Projection(projection) => {
                        let assoc_data = self.db.associated_ty_data(projection.associated_ty_id);
                        if !fn_traits.contains(&assoc_data.trait_id) {
                            return None;
                        }

                        // Skip `Self`, get the type argument.
                        let arg = projection.substitution.as_slice(&Interner).get(1)?;
                        match arg.ty(&Interner)?.kind(&Interner) {
                            TyKind::Tuple(_, subst) => {
                                let generic_args = subst.as_slice(&Interner);
                                let mut sig_tys = Vec::new();
                                for arg in generic_args {
                                    sig_tys.push(arg.ty(&Interner)?.clone());
                                }
                                sig_tys.push(eq.ty.clone());

                                cov_mark::hit!(dyn_fn_param_informs_call_site_closure_signature);
                                return Some(FnPointer {
                                    num_binders: 0,
                                    sig: FnSig {
                                        abi: (),
                                        safety: chalk_ir::Safety::Safe,
                                        variadic: false,
                                    },
                                    substitution: FnSubst(
                                        Substitution::from_iter(&Interner, sig_tys.clone())
                                            .shifted_in(&Interner),
                                    ),
                                });
                            }
                            _ => {}
                        }
                    }
                    AliasTy::Opaque(_) => {}
                },
                _ => {}
            }
        }

        None
    }
}
<|endoftext|>