<commit_before>use crate::backend::UsesAnsiSavepointSyntax;
use crate::connection::{Connection, SimpleConnection};
use crate::result::{DatabaseErrorKind, Error, QueryResult};

/// Manages the internal transaction state for a connection.
///
/// You will not need to interact with this trait, unless you are writing an
/// implementation of [`Connection`](trait.Connection.html).
pub trait TransactionManager<Conn: Connection> {
    /// Begin a new transaction or savepoint
    ///
    /// If the transaction depth is greater than 0,
    /// this should create a savepoint instead.
    /// This function is expected to increment the transaction depth by 1.
    fn begin_transaction(&self, conn: &Conn) -> QueryResult<()>;

    /// Rollback the inner-most transaction or savepoint
    ///
    /// If the transaction depth is greater than 1,
    /// this should rollback to the most recent savepoint.
    /// This function is expected to decrement the transaction depth by 1.
    fn rollback_transaction(&self, conn: &Conn) -> QueryResult<()>;

    /// Commit the inner-most transaction or savepoint
    ///
    /// If the transaction depth is greater than 1,
    /// this should release the most recent savepoint.
    /// This function is expected to decrement the transaction depth by 1.
    fn commit_transaction(&self, conn: &Conn) -> QueryResult<()>;

    /// Fetch the current transaction depth
    ///
    /// Used to ensure that `begin_test_transaction` is not called when already
    /// inside of a transaction.
    fn get_transaction_depth(&self) -> u32;
}

use std::cell::Cell;

/// An implementation of `TransactionManager` which can be used for backends
/// which use ANSI standard syntax for savepoints such as SQLite and PostgreSQL.
#[allow(missing_debug_implementations)]
#[derive(Default)]
pub struct AnsiTransactionManager {
    transaction_depth: Cell<i32>,
}

impl AnsiTransactionManager {
    /// Create a new transaction manager
    pub fn new() -> Self {
        AnsiTransactionManager::default()
    }

    fn change_transaction_depth(&self, by: i32, query: QueryResult<()>) -> QueryResult<()> {
        if query.is_ok() {
            self.transaction_depth
                .set(self.transaction_depth.get() + by)
        }
        query
    }

    /// Begin a transaction with custom SQL
    ///
    /// This is used by connections to implement more complex transaction APIs
    /// to set things such as isolation levels.
    /// Returns an error if already inside of a transaction.
    pub fn begin_transaction_sql<Conn>(&self, conn: &Conn, sql: &str) -> QueryResult<()>
    where
        Conn: SimpleConnection,
    {
        use crate::result::Error::AlreadyInTransaction;

        if self.transaction_depth.get() == 0 {
            self.change_transaction_depth(1, conn.batch_execute(sql))
        } else {
            Err(AlreadyInTransaction)
        }
    }
}

impl<Conn> TransactionManager<Conn> for AnsiTransactionManager
where
    Conn: Connection,
    Conn::Backend: UsesAnsiSavepointSyntax,
{
    fn begin_transaction(&self, conn: &Conn) -> QueryResult<()> {
        let transaction_depth = self.transaction_depth.get();
        self.change_transaction_depth(
            1,
            if transaction_depth == 0 {
                conn.batch_execute("BEGIN")
            } else {
                conn.batch_execute(&format!("SAVEPOINT diesel_savepoint_{}", transaction_depth))
            },
        )
    }

    fn rollback_transaction(&self, conn: &Conn) -> QueryResult<()> {
        let transaction_depth = self.transaction_depth.get();
        self.change_transaction_depth(
            -1,
            if transaction_depth == 1 {
                conn.batch_execute("ROLLBACK")
            } else {
                conn.batch_execute(&format!(
                    "ROLLBACK TO SAVEPOINT diesel_savepoint_{}",
                    transaction_depth - 1
                ))
            },
        )
    }

    /// If the transaction fails to commit due to a `SerializationFailure` or a
    /// `ReadOnlyTransaction` rollback will be attempted. If it succeeds,
    /// original error will be returned. If it fails, rollback error will be returned.
    fn commit_transaction(&self, conn: &Conn) -> QueryResult<()> {
        let transaction_depth = self.transaction_depth.get();
        let commit_result = self.change_transaction_depth(
            -1,
            if transaction_depth <= 1 {
                conn.batch_execute("COMMIT")
            } else {
                conn.batch_execute(&format!(
                    "RELEASE SAVEPOINT diesel_savepoint_{}",
                    transaction_depth - 1
                ))
            },
        );
        if let Err(ref err) = commit_result {
            match err {
                Error::DatabaseError(DatabaseErrorKind::SerializationFailure, _)
                | Error::DatabaseError(DatabaseErrorKind::ReadOnlyTransaction, _) => {
                    self.rollback_transaction(conn)?;
                }
                _ => {}
            }
        }
        commit_result
    }

    fn get_transaction_depth(&self) -> u32 {
        self.transaction_depth.get() as u32
    }
}
<commit_msg>Only catch on `COMMIT`, not on `SAVEPOINT`<commit_after>use crate::backend::UsesAnsiSavepointSyntax;
use crate::connection::{Connection, SimpleConnection};
use crate::result::{DatabaseErrorKind, Error, QueryResult};

/// Manages the internal transaction state for a connection.
///
/// You will not need to interact with this trait, unless you are writing an
/// implementation of [`Connection`](trait.Connection.html).
pub trait TransactionManager<Conn: Connection> {
    /// Begin a new transaction or savepoint
    ///
    /// If the transaction depth is greater than 0,
    /// this should create a savepoint instead.
    /// This function is expected to increment the transaction depth by 1.
    fn begin_transaction(&self, conn: &Conn) -> QueryResult<()>;

    /// Rollback the inner-most transaction or savepoint
    ///
    /// If the transaction depth is greater than 1,
    /// this should rollback to the most recent savepoint.
    /// This function is expected to decrement the transaction depth by 1.
    fn rollback_transaction(&self, conn: &Conn) -> QueryResult<()>;

    /// Commit the inner-most transaction or savepoint
    ///
    /// If the transaction depth is greater than 1,
    /// this should release the most recent savepoint.
    /// This function is expected to decrement the transaction depth by 1.
    fn commit_transaction(&self, conn: &Conn) -> QueryResult<()>;

    /// Fetch the current transaction depth
    ///
    /// Used to ensure that `begin_test_transaction` is not called when already
    /// inside of a transaction.
    fn get_transaction_depth(&self) -> u32;
}

use std::cell::Cell;

/// An implementation of `TransactionManager` which can be used for backends
/// which use ANSI standard syntax for savepoints such as SQLite and PostgreSQL.
#[allow(missing_debug_implementations)]
#[derive(Default)]
pub struct AnsiTransactionManager {
    transaction_depth: Cell<i32>,
}

impl AnsiTransactionManager {
    /// Create a new transaction manager
    pub fn new() -> Self {
        AnsiTransactionManager::default()
    }

    fn change_transaction_depth(&self, by: i32, query: QueryResult<()>) -> QueryResult<()> {
        if query.is_ok() {
            self.transaction_depth
                .set(self.transaction_depth.get() + by)
        }
        query
    }

    /// Begin a transaction with custom SQL
    ///
    /// This is used by connections to implement more complex transaction APIs
    /// to set things such as isolation levels.
    /// Returns an error if already inside of a transaction.
    pub fn begin_transaction_sql<Conn>(&self, conn: &Conn, sql: &str) -> QueryResult<()>
    where
        Conn: SimpleConnection,
    {
        use crate::result::Error::AlreadyInTransaction;

        if self.transaction_depth.get() == 0 {
            self.change_transaction_depth(1, conn.batch_execute(sql))
        } else {
            Err(AlreadyInTransaction)
        }
    }
}

impl<Conn> TransactionManager<Conn> for AnsiTransactionManager
where
    Conn: Connection,
    Conn::Backend: UsesAnsiSavepointSyntax,
{
    fn begin_transaction(&self, conn: &Conn) -> QueryResult<()> {
        let transaction_depth = self.transaction_depth.get();
        self.change_transaction_depth(
            1,
            if transaction_depth == 0 {
                conn.batch_execute("BEGIN")
            } else {
                conn.batch_execute(&format!("SAVEPOINT diesel_savepoint_{}", transaction_depth))
            },
        )
    }

    fn rollback_transaction(&self, conn: &Conn) -> QueryResult<()> {
        let transaction_depth = self.transaction_depth.get();
        self.change_transaction_depth(
            -1,
            if transaction_depth == 1 {
                conn.batch_execute("ROLLBACK")
            } else {
                conn.batch_execute(&format!(
                    "ROLLBACK TO SAVEPOINT diesel_savepoint_{}",
                    transaction_depth - 1
                ))
            },
        )
    }

    /// If the transaction fails to commit due to a `SerializationFailure` or a
    /// `ReadOnlyTransaction` rollback will be attempted. If it succeeds,
    /// original error will be returned. If it fails, rollback error will be returned.
    fn commit_transaction(&self, conn: &Conn) -> QueryResult<()> {
        let transaction_depth = self.transaction_depth.get();
        self.change_transaction_depth(
            -1,
            if transaction_depth <= 1 {
                conn.batch_execute("COMMIT").or_else(|err| match err {
                    Error::DatabaseError(DatabaseErrorKind::SerializationFailure, _)
                    | Error::DatabaseError(DatabaseErrorKind::ReadOnlyTransaction, _) => {
                        conn.batch_execute("ROLLBACK").and(Err(err))
                    }
                    other_err => Err(other_err),
                })
            } else {
                conn.batch_execute(&format!(
                    "RELEASE SAVEPOINT diesel_savepoint_{}",
                    transaction_depth - 1
                ))
            },
        )
    }

    fn get_transaction_depth(&self) -> u32 {
        self.transaction_depth.get() as u32
    }
}
<|endoftext|>


<commit_before>use crate::{sys, Token};

use std::fmt;

/// A readiness event.
///
/// `Event` is a readiness state paired with a [`Token`]. It is returned by
/// [`Poll::poll`].
///
/// For more documentation on polling and events, see [`Poll`].
///
/// [`Poll::poll`]: ../struct.Poll.html#method.poll
/// [`Poll`]: ../struct.Poll.html
/// [`Token`]: ../struct.Token.html
#[repr(transparent)]
pub struct Event {
    inner: sys::Event,
}

impl Event {
    /// Returns the event's token.
    pub fn token(&self) -> Token {
        sys::event::token(&self.inner)
    }

    /// Returns true if the event contains readable readiness.
    pub fn is_readable(&self) -> bool {
        sys::event::is_readable(&self.inner)
    }

    /// Returns true if the event contains writable readiness.
    pub fn is_writable(&self) -> bool {
        sys::event::is_writable(&self.inner)
    }

    /// Returns true if the event contains error readiness.
    ///
    /// Error events occur when the socket enters an error state. In this case,
    /// the socket will also receive a readable or writable event. Reading or
    /// writing to the socket will result in an error.
    ///
    /// # Notes
    ///
    /// Method is available on all platforms, but not all platforms trigger the
    /// error event.
    ///
    /// The table below shows what flags are checked on what OS.
    ///
    /// | [OS selector] | Flag(s) checked |
    /// |---------------|-----------------|
    /// | [epoll]       | `EPOLLERR`      |
    /// | [kqueue]      | `EV_ERROR` and `EV_EOF` with `fflags` set to `0`. |
    ///
    /// [OS selector]: ../struct.Poll.html#implementation-notes
    /// [epoll]: http://man7.org/linux/man-pages/man7/epoll.7.html
    /// [kqueue]: https://www.freebsd.org/cgi/man.cgi?query=kqueue&sektion=2
    pub fn is_error(&self) -> bool {
        sys::event::is_error(&self.inner)
    }

    /// Returns true if the event contains read closed readiness.
    ///
    /// # Notes
    ///
    /// Read closed readiness can be expected after any of the following have
    /// occurred:
    /// * The local stream has shutdown the read half of its socket
    /// * The local stream has shtudown both the read half and the write half
    ///   of its socket
    /// * The peer stream has shtudown the write half its socket; this sends a
    ///   `FIN` packet that has been received by the local stream
    ///
    /// Method is a best effort implementation. While some platforms may not
    /// return readiness when read half is closed, it is guaranteed that
    /// false-positives will not occur.
    ///
    /// The table below shows what flags are checked on what OS.
    ///
    /// | [OS selector] | Flag(s) checked |
    /// |---------------|-----------------|
    /// | [epoll]       | `EPOLLHUP`, or  |
    /// |               | `EPOLLIN` and EPOLLRDHUP` |
    /// | [kqueue]      | `EV_EOF`        |
    ///
    /// [OS selector]: ../struct.Poll.html#implementation-notes
    /// [epoll]: http://man7.org/linux/man-pages/man7/epoll.7.html
    /// [kqueue]: https://www.freebsd.org/cgi/man.cgi?query=kqueue&sektion=2
    pub fn is_read_closed(&self) -> bool {
        sys::event::is_read_closed(&self.inner)
    }

    /// Returns true if the event contains write closed readiness.
    ///
    /// # Notes
    ///
    /// On [epoll] this is essentially a check for `EPOLLHUP` flag as the
    /// local stream shutting down its write half does not trigger this event.
    ///
    /// On [kqueue] the local stream shutting down the write half of its
    /// socket will trigger this event.
    ///
    /// Method is a best effort implementation. While some platforms may not
    /// return readiness when write half is closed, it is guaranteed that
    /// false-positives will not occur.
    ///
    /// The table below shows what flags are checked on what OS.
    ///
    /// | [OS selector] | Flag(s) checked |
    /// |---------------|-----------------|
    /// | [epoll]       | `EPOLLHUP`, or  |
    /// |               | `EPOLLOUT` and `EPOLLERR` |
    /// | [kqueue]      | `EV_EOF`        |
    ///
    /// [OS selector]: ../struct.Poll.html#implementation-notes
    /// [epoll]: http://man7.org/linux/man-pages/man7/epoll.7.html
    /// [kqueue]: https://www.freebsd.org/cgi/man.cgi?query=kqueue&sektion=2
    pub fn is_write_closed(&self) -> bool {
        sys::event::is_write_closed(&self.inner)
    }

    /// Returns true if the event contains priority readiness.
    ///
    /// # Notes
    ///
    /// Method is available on all platforms, but not all platforms trigger the
    /// priority event.
    ///
    /// The table below shows what flags are checked on what OS.
    ///
    /// | [OS selector] | Flag(s) checked |
    /// |---------------|-----------------|
    /// | [epoll]       | `EPOLLPRI`      |
    /// | [kqueue]      | *Not supported* |
    ///
    /// [OS selector]: ../struct.Poll.html#implementation-notes
    /// [epoll]: http://man7.org/linux/man-pages/man7/epoll.7.html
    /// [kqueue]: https://www.freebsd.org/cgi/man.cgi?query=kqueue&sektion=2
    #[inline]
    pub fn is_priority(&self) -> bool {
        sys::event::is_priority(&self.inner)
    }

    /// Returns true if the event contains AIO readiness.
    ///
    /// # Notes
    ///
    /// Method is available on all platforms, but not all platforms support AIO.
    ///
    /// The table below shows what flags are checked on what OS.
    ///
    /// | [OS selector] | Flag(s) checked |
    /// |---------------|-----------------|
    /// | [epoll]       | *Not supported* |
    /// | [kqueue]1 | `EVFILT_AIO` |
    ///
    /// 1: Only supported on DragonFly BSD, FreeBSD, iOS and macOS.
    ///
    /// [OS selector]: ../struct.Poll.html#implementation-notes
    /// [epoll]: http://man7.org/linux/man-pages/man7/epoll.7.html
    /// [kqueue]: https://www.freebsd.org/cgi/man.cgi?query=kqueue&sektion=2
    pub fn is_aio(&self) -> bool {
        sys::event::is_aio(&self.inner)
    }

    /// Returns true if the event contains LIO readiness.
    ///
    /// # Notes
    ///
    /// Method is available on all platforms, but only FreeBSD supports LIO. On
    /// FreeBSD this method checks the `EVFILT_LIO` flag.
    pub fn is_lio(&self) -> bool {
        sys::event::is_lio(&self.inner)
    }

    /// Create a reference to an `Event` from a platform specific event.
    pub(crate) fn from_sys_event_ref(sys_event: &sys::Event) -> &Event {
        unsafe {
            // This is safe because the memory layout of `Event` is
            // the same as `sys::Event` due to the `repr(transparent)` attribute.
            &*(sys_event as *const sys::Event as *const Event)
        }
    }
}

/// When the [alternate] flag is enabled this will print platform specific
/// details, for example the fields of the `kevent` structure on platforms that
/// use `kqueue(2)`. Note however that the output of this implementation is
/// **not** consider a part of the stable API.
///
/// [alternate]: fmt::Formatter::alternate
impl fmt::Debug for Event {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let alternate = f.alternate();
        let mut d = f.debug_struct("Event");
        d.field("token", &self.token())
            .field("readable", &self.is_readable())
            .field("writable", &self.is_writable())
            .field("error", &self.is_error())
            .field("read_closed", &self.is_read_closed())
            .field("write_closed", &self.is_write_closed())
            .field("priority", &self.is_priority())
            .field("aio", &self.is_aio())
            .field("lio", &self.is_lio());

        if alternate {
            struct EventDetails<'a>(&'a sys::Event);

            impl<'a> fmt::Debug for EventDetails<'a> {
                fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                    sys::event::debug_details(f, self.0)
                }
            }

            d.field("details", &EventDetails(&self.inner)).finish()
        } else {
            d.finish()
        }
    }
}
<commit_msg>Fix missing qoute<commit_after>use crate::{sys, Token};

use std::fmt;

/// A readiness event.
///
/// `Event` is a readiness state paired with a [`Token`]. It is returned by
/// [`Poll::poll`].
///
/// For more documentation on polling and events, see [`Poll`].
///
/// [`Poll::poll`]: ../struct.Poll.html#method.poll
/// [`Poll`]: ../struct.Poll.html
/// [`Token`]: ../struct.Token.html
#[repr(transparent)]
pub struct Event {
    inner: sys::Event,
}

impl Event {
    /// Returns the event's token.
    pub fn token(&self) -> Token {
        sys::event::token(&self.inner)
    }

    /// Returns true if the event contains readable readiness.
    pub fn is_readable(&self) -> bool {
        sys::event::is_readable(&self.inner)
    }

    /// Returns true if the event contains writable readiness.
    pub fn is_writable(&self) -> bool {
        sys::event::is_writable(&self.inner)
    }

    /// Returns true if the event contains error readiness.
    ///
    /// Error events occur when the socket enters an error state. In this case,
    /// the socket will also receive a readable or writable event. Reading or
    /// writing to the socket will result in an error.
    ///
    /// # Notes
    ///
    /// Method is available on all platforms, but not all platforms trigger the
    /// error event.
    ///
    /// The table below shows what flags are checked on what OS.
    ///
    /// | [OS selector] | Flag(s) checked |
    /// |---------------|-----------------|
    /// | [epoll]       | `EPOLLERR`      |
    /// | [kqueue]      | `EV_ERROR` and `EV_EOF` with `fflags` set to `0`. |
    ///
    /// [OS selector]: ../struct.Poll.html#implementation-notes
    /// [epoll]: http://man7.org/linux/man-pages/man7/epoll.7.html
    /// [kqueue]: https://www.freebsd.org/cgi/man.cgi?query=kqueue&sektion=2
    pub fn is_error(&self) -> bool {
        sys::event::is_error(&self.inner)
    }

    /// Returns true if the event contains read closed readiness.
    ///
    /// # Notes
    ///
    /// Read closed readiness can be expected after any of the following have
    /// occurred:
    /// * The local stream has shutdown the read half of its socket
    /// * The local stream has shtudown both the read half and the write half
    ///   of its socket
    /// * The peer stream has shtudown the write half its socket; this sends a
    ///   `FIN` packet that has been received by the local stream
    ///
    /// Method is a best effort implementation. While some platforms may not
    /// return readiness when read half is closed, it is guaranteed that
    /// false-positives will not occur.
    ///
    /// The table below shows what flags are checked on what OS.
    ///
    /// | [OS selector] | Flag(s) checked |
    /// |---------------|-----------------|
    /// | [epoll]       | `EPOLLHUP`, or  |
    /// |               | `EPOLLIN` and `EPOLLRDHUP` |
    /// | [kqueue]      | `EV_EOF`        |
    ///
    /// [OS selector]: ../struct.Poll.html#implementation-notes
    /// [epoll]: http://man7.org/linux/man-pages/man7/epoll.7.html
    /// [kqueue]: https://www.freebsd.org/cgi/man.cgi?query=kqueue&sektion=2
    pub fn is_read_closed(&self) -> bool {
        sys::event::is_read_closed(&self.inner)
    }

    /// Returns true if the event contains write closed readiness.
    ///
    /// # Notes
    ///
    /// On [epoll] this is essentially a check for `EPOLLHUP` flag as the
    /// local stream shutting down its write half does not trigger this event.
    ///
    /// On [kqueue] the local stream shutting down the write half of its
    /// socket will trigger this event.
    ///
    /// Method is a best effort implementation. While some platforms may not
    /// return readiness when write half is closed, it is guaranteed that
    /// false-positives will not occur.
    ///
    /// The table below shows what flags are checked on what OS.
    ///
    /// | [OS selector] | Flag(s) checked |
    /// |---------------|-----------------|
    /// | [epoll]       | `EPOLLHUP`, or  |
    /// |               | `EPOLLOUT` and `EPOLLERR` |
    /// | [kqueue]      | `EV_EOF`        |
    ///
    /// [OS selector]: ../struct.Poll.html#implementation-notes
    /// [epoll]: http://man7.org/linux/man-pages/man7/epoll.7.html
    /// [kqueue]: https://www.freebsd.org/cgi/man.cgi?query=kqueue&sektion=2
    pub fn is_write_closed(&self) -> bool {
        sys::event::is_write_closed(&self.inner)
    }

    /// Returns true if the event contains priority readiness.
    ///
    /// # Notes
    ///
    /// Method is available on all platforms, but not all platforms trigger the
    /// priority event.
    ///
    /// The table below shows what flags are checked on what OS.
    ///
    /// | [OS selector] | Flag(s) checked |
    /// |---------------|-----------------|
    /// | [epoll]       | `EPOLLPRI`      |
    /// | [kqueue]      | *Not supported* |
    ///
    /// [OS selector]: ../struct.Poll.html#implementation-notes
    /// [epoll]: http://man7.org/linux/man-pages/man7/epoll.7.html
    /// [kqueue]: https://www.freebsd.org/cgi/man.cgi?query=kqueue&sektion=2
    #[inline]
    pub fn is_priority(&self) -> bool {
        sys::event::is_priority(&self.inner)
    }

    /// Returns true if the event contains AIO readiness.
    ///
    /// # Notes
    ///
    /// Method is available on all platforms, but not all platforms support AIO.
    ///
    /// The table below shows what flags are checked on what OS.
    ///
    /// | [OS selector] | Flag(s) checked |
    /// |---------------|-----------------|
    /// | [epoll]       | *Not supported* |
    /// | [kqueue]1 | `EVFILT_AIO` |
    ///
    /// 1: Only supported on DragonFly BSD, FreeBSD, iOS and macOS.
    ///
    /// [OS selector]: ../struct.Poll.html#implementation-notes
    /// [epoll]: http://man7.org/linux/man-pages/man7/epoll.7.html
    /// [kqueue]: https://www.freebsd.org/cgi/man.cgi?query=kqueue&sektion=2
    pub fn is_aio(&self) -> bool {
        sys::event::is_aio(&self.inner)
    }

    /// Returns true if the event contains LIO readiness.
    ///
    /// # Notes
    ///
    /// Method is available on all platforms, but only FreeBSD supports LIO. On
    /// FreeBSD this method checks the `EVFILT_LIO` flag.
    pub fn is_lio(&self) -> bool {
        sys::event::is_lio(&self.inner)
    }

    /// Create a reference to an `Event` from a platform specific event.
    pub(crate) fn from_sys_event_ref(sys_event: &sys::Event) -> &Event {
        unsafe {
            // This is safe because the memory layout of `Event` is
            // the same as `sys::Event` due to the `repr(transparent)` attribute.
            &*(sys_event as *const sys::Event as *const Event)
        }
    }
}

/// When the [alternate] flag is enabled this will print platform specific
/// details, for example the fields of the `kevent` structure on platforms that
/// use `kqueue(2)`. Note however that the output of this implementation is
/// **not** consider a part of the stable API.
///
/// [alternate]: fmt::Formatter::alternate
impl fmt::Debug for Event {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let alternate = f.alternate();
        let mut d = f.debug_struct("Event");
        d.field("token", &self.token())
            .field("readable", &self.is_readable())
            .field("writable", &self.is_writable())
            .field("error", &self.is_error())
            .field("read_closed", &self.is_read_closed())
            .field("write_closed", &self.is_write_closed())
            .field("priority", &self.is_priority())
            .field("aio", &self.is_aio())
            .field("lio", &self.is_lio());

        if alternate {
            struct EventDetails<'a>(&'a sys::Event);

            impl<'a> fmt::Debug for EventDetails<'a> {
                fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                    sys::event::debug_details(f, self.0)
                }
            }

            d.field("details", &EventDetails(&self.inner)).finish()
        } else {
            d.finish()
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
// Copyright 2013 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The normal and derived distributions.

use crate::utils::ziggurat;
use num_traits::Float;
use crate::{ziggurat_tables, Distribution, Open01};
use rand::Rng;
use core::fmt;

/// Samples floating-point numbers according to the normal distribution
/// `N(0, 1)` (a.k.a. a standard normal, or Gaussian). This is equivalent to
/// `Normal::new(0.0, 1.0)` but faster.
///
/// See `Normal` for the general normal distribution.
///
/// Implemented via the ZIGNOR variant[^1] of the Ziggurat method.
///
/// [^1]: Jurgen A. Doornik (2005). [*An Improved Ziggurat Method to
///       Generate Normal Random Samples*](
///       https://www.doornik.com/research/ziggurat.pdf).
///       Nuffield College, Oxford
///
/// # Example
/// ```
/// use rand::prelude::*;
/// use rand_distr::StandardNormal;
///
/// let val: f64 = thread_rng().sample(StandardNormal);
/// println!("{}", val);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct StandardNormal;

impl Distribution<f32> for StandardNormal {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f32 {
        // TODO: use optimal 32-bit implementation
        let x: f64 = self.sample(rng);
        x as f32
    }
}

impl Distribution<f64> for StandardNormal {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
        #[inline]
        fn pdf(x: f64) -> f64 {
            (-x * x / 2.0).exp()
        }
        #[inline]
        fn zero_case<R: Rng + ?Sized>(rng: &mut R, u: f64) -> f64 {
            // compute a random number in the tail by hand

            // strange initial conditions, because the loop is not
            // do-while, so the condition should be true on the first
            // run, they get overwritten anyway (0 < 1, so these are
            // good).
            let mut x = 1.0f64;
            let mut y = 0.0f64;

            while -2.0 * y < x * x {
                let x_: f64 = rng.sample(Open01);
                let y_: f64 = rng.sample(Open01);

                x = x_.ln() / ziggurat_tables::ZIG_NORM_R;
                y = y_.ln();
            }

            if u < 0.0 {
                x - ziggurat_tables::ZIG_NORM_R
            } else {
                ziggurat_tables::ZIG_NORM_R - x
            }
        }

        ziggurat(
            rng,
            true, // this is symmetric
            &ziggurat_tables::ZIG_NORM_X,
            &ziggurat_tables::ZIG_NORM_F,
            pdf,
            zero_case,
        )
    }
}

/// The normal distribution `N(mean, std_dev**2)`.
///
/// This uses the ZIGNOR variant of the Ziggurat method, see [`StandardNormal`]
/// for more details.
///
/// Note that [`StandardNormal`] is an optimised implementation for mean 0, and
/// standard deviation 1.
///
/// # Example
///
/// ```
/// use rand_distr::{Normal, Distribution};
///
/// // mean 2, standard deviation 3
/// let normal = Normal::new(2.0, 3.0).unwrap();
/// let v = normal.sample(&mut rand::thread_rng());
/// println!("{} is from a N(2, 9) distribution", v)
/// ```
///
/// [`StandardNormal`]: crate::StandardNormal
#[derive(Clone, Copy, Debug)]
pub struct Normal<F>
where F: Float, StandardNormal: Distribution<F>
{
    mean: F,
    std_dev: F,
}

/// Error type returned from `Normal::new` and `LogNormal::new`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    /// The mean value is too small (log-normal samples must be positive)
    MeanTooSmall,
    /// The standard deviation or other dispersion parameter is negative,
    /// infinite, or not-a-number.
    BadVariance,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            Error::MeanTooSmall => "mean < 0 or NaN in log-normal distribution",
            Error::BadVariance => "variation parameter is negative or non-finite in (log)normal distribution",
        })
    }
}

#[cfg(feature = "std")]
impl std::error::Error for Error {}

impl<F> Normal<F>
where F: Float, StandardNormal: Distribution<F>
{
    /// Construct, from mean and standard deviation
    ///
    /// Parameters:
    ///
    /// -   mean (`μ`, unrestricted)
    /// -   standard deviation (`σ`, must be finite)
    #[inline]
    pub fn new(mean: F, std_dev: F) -> Result<Normal<F>, Error> {
        if !std_dev.is_finite() {
            return Err(Error::BadVariance);
        }
        Ok(Normal { mean, std_dev })
    }

    /// Construct, from mean and coefficient of variation
    ///
    /// Parameters:
    ///
    /// -   mean (`μ`, unrestricted)
    /// -   coefficient of variation (`cv = abs(σ / μ)`)
    #[inline]
    pub fn from_mean_cv(mean: F, cv: F) -> Result<Normal<F>, Error> {
        if !cv.is_finite() || cv < F::zero() {
            return Err(Error::BadVariance);
        }
        let std_dev = cv * mean;
        Ok(Normal { mean, std_dev })
    }

    /// Sample from a z-score
    ///
    /// This may be useful for generating correlated samples, as follows.
    /// ```
    /// # use rand::prelude::*;
    /// # use rand_distr::{Normal, StandardNormal};
    /// let mut rng = thread_rng();
    /// let z = StandardNormal.sample(&mut rng);
    /// let x1 = Normal::new(0.0, 1.0).unwrap().from_zscore(z);
    /// let x2 = Normal::new(2.0, -1.0).unwrap().from_zscore(z);
    /// ```
    #[inline]
    pub fn from_zscore(&self, zscore: F) -> F {
        self.mean + self.std_dev * zscore
    }
}

impl<F> Distribution<F> for Normal<F>
where F: Float, StandardNormal: Distribution<F>
{
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
        self.from_zscore(rng.sample(StandardNormal))
    }
}

/// The log-normal distribution `ln N(mean, std_dev**2)`.
///
/// If `X` is log-normal distributed, then `ln(X)` is `N(mean, std_dev**2)`
/// distributed.
///
/// # Example
///
/// ```
/// use rand_distr::{LogNormal, Distribution};
///
/// // mean 2, standard deviation 3
/// let log_normal = LogNormal::new(2.0, 3.0).unwrap();
/// let v = log_normal.sample(&mut rand::thread_rng());
/// println!("{} is from an ln N(2, 9) distribution", v)
/// ```
#[derive(Clone, Copy, Debug)]
pub struct LogNormal<F>
where F: Float, StandardNormal: Distribution<F>
{
    norm: Normal<F>,
}

impl<F> LogNormal<F>
where F: Float, StandardNormal: Distribution<F>
{
    /// Construct, from (log-space) mean and standard deviation
    ///
    /// Parameters are the "standard" log-space measures (these are the mean
    /// and standard deviation of the logarithm of samples):
    ///
    /// -   `mu` (`μ`, unrestricted) is the mean of the underlying distribution
    /// -   `sigma` (`σ`, must be finite) is the standard deviation of the
    ///     underlying Normal distribution
    #[inline]
    pub fn new(mu: F, sigma: F) -> Result<LogNormal<F>, Error> {
        let norm = Normal::new(mu, sigma)?;
        Ok(LogNormal { norm })
    }

    /// Construct, from (linear-space) mean and coefficient of variation
    ///
    /// Parameters are linear-space measures:
    ///
    /// -   mean (`μ > 0`) is the (real) mean of the distribution
    /// -   coefficient of variation (`cv = σ / μ`, requiring `cv ≥ 0`) is a
    ///     standardized measure of dispersion
    ///
    /// As a special exception, `μ = 0, cv = 0` is allowed (samples are `-inf`).
    #[inline]
    pub fn from_mean_cv(mean: F, cv: F) -> Result<LogNormal<F>, Error> {
        if cv == F::zero() {
            let mu = mean.ln();
            let norm = Normal::new(mu, F::zero()).unwrap();
            return Ok(LogNormal { norm });
        }
        if !(mean > F::zero()) {
            return Err(Error::MeanTooSmall);
        }
        if !(cv >= F::zero()) {
            return Err(Error::BadVariance);
        }

        // Using X ~ lognormal(μ, σ), CV² = Var(X) / E(X)²
        // E(X) = exp(μ + σ² / 2) = exp(μ) × exp(σ² / 2)
        // Var(X) = exp(2μ + σ²)(exp(σ²) - 1) = E(X)² × (exp(σ²) - 1)
        // but Var(X) = (CV × E(X))² so CV² = exp(σ²) - 1
        // thus σ² = log(CV² + 1)
        // and exp(μ) = E(X) / exp(σ² / 2) = E(X) / sqrt(CV² + 1)
        let a = F::one() + cv * cv; // e
        let mu = F::from(0.5).unwrap() * (mean * mean / a).ln();
        let sigma = a.ln().sqrt();
        let norm = Normal::new(mu, sigma)?;
        Ok(LogNormal { norm })
    }

    /// Sample from a z-score
    ///
    /// This may be useful for generating correlated samples, as follows.
    /// ```
    /// # use rand::prelude::*;
    /// # use rand_distr::{LogNormal, StandardNormal};
    /// let mut rng = thread_rng();
    /// let z = StandardNormal.sample(&mut rng);
    /// let x1 = LogNormal::from_mean_cv(3.0, 1.0).unwrap().from_zscore(z);
    /// let x2 = LogNormal::from_mean_cv(2.0, 1.0).unwrap().from_zscore(z);
    /// ```
    #[inline]
    pub fn from_zscore(&self, zscore: F) -> F {
        self.norm.from_zscore(zscore).exp()
    }
}

impl<F> Distribution<F> for LogNormal<F>
where F: Float, StandardNormal: Distribution<F>
{
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
        self.norm.sample(rng).exp()
    }
}

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

    #[test]
    fn test_normal() {
        let norm = Normal::new(10.0, 10.0).unwrap();
        let mut rng = crate::test::rng(210);
        for _ in 0..1000 {
            norm.sample(&mut rng);
        }
    }
    #[test]
    fn test_normal_cv() {
        let norm = Normal::from_mean_cv(1024.0, 1.0 / 256.0).unwrap();
        assert_eq!((norm.mean, norm.std_dev), (1024.0, 4.0));
    }
    #[test]
    fn test_normal_invalid_sd() {
        assert!(Normal::from_mean_cv(10.0, -1.0).is_err());
    }

    #[test]
    fn test_log_normal() {
        let lnorm = LogNormal::new(10.0, 10.0).unwrap();
        let mut rng = crate::test::rng(211);
        for _ in 0..1000 {
            lnorm.sample(&mut rng);
        }
    }
    #[test]
    fn test_log_normal_cv() {
        let lnorm = LogNormal::from_mean_cv(0.0, 0.0).unwrap();
        assert_eq!((lnorm.norm.mean, lnorm.norm.std_dev), (-std::f64::INFINITY, 0.0));

        let lnorm = LogNormal::from_mean_cv(1.0, 0.0).unwrap();
        assert_eq!((lnorm.norm.mean, lnorm.norm.std_dev), (0.0, 0.0));

        let e = std::f64::consts::E;
        let lnorm = LogNormal::from_mean_cv(e.sqrt(), (e - 1.0).sqrt()).unwrap();
        assert_almost_eq!(lnorm.norm.mean, 0.0, 2e-16);
        assert_almost_eq!(lnorm.norm.std_dev, 1.0, 2e-16);

        let lnorm = LogNormal::from_mean_cv(e.powf(1.5), (e - 1.0).sqrt()).unwrap();
        assert_eq!((lnorm.norm.mean, lnorm.norm.std_dev), (1.0, 1.0));
    }
    #[test]
    fn test_log_normal_invalid_sd() {
        assert!(LogNormal::from_mean_cv(-1.0, 1.0).is_err());
        assert!(LogNormal::from_mean_cv(0.0, 1.0).is_err());
        assert!(LogNormal::from_mean_cv(1.0, -1.0).is_err());
    }
}
<commit_msg>Update correlated normal samples example<commit_after>// Copyright 2018 Developers of the Rand project.
// Copyright 2013 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The normal and derived distributions.

use crate::utils::ziggurat;
use num_traits::Float;
use crate::{ziggurat_tables, Distribution, Open01};
use rand::Rng;
use core::fmt;

/// Samples floating-point numbers according to the normal distribution
/// `N(0, 1)` (a.k.a. a standard normal, or Gaussian). This is equivalent to
/// `Normal::new(0.0, 1.0)` but faster.
///
/// See `Normal` for the general normal distribution.
///
/// Implemented via the ZIGNOR variant[^1] of the Ziggurat method.
///
/// [^1]: Jurgen A. Doornik (2005). [*An Improved Ziggurat Method to
///       Generate Normal Random Samples*](
///       https://www.doornik.com/research/ziggurat.pdf).
///       Nuffield College, Oxford
///
/// # Example
/// ```
/// use rand::prelude::*;
/// use rand_distr::StandardNormal;
///
/// let val: f64 = thread_rng().sample(StandardNormal);
/// println!("{}", val);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct StandardNormal;

impl Distribution<f32> for StandardNormal {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f32 {
        // TODO: use optimal 32-bit implementation
        let x: f64 = self.sample(rng);
        x as f32
    }
}

impl Distribution<f64> for StandardNormal {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f64 {
        #[inline]
        fn pdf(x: f64) -> f64 {
            (-x * x / 2.0).exp()
        }
        #[inline]
        fn zero_case<R: Rng + ?Sized>(rng: &mut R, u: f64) -> f64 {
            // compute a random number in the tail by hand

            // strange initial conditions, because the loop is not
            // do-while, so the condition should be true on the first
            // run, they get overwritten anyway (0 < 1, so these are
            // good).
            let mut x = 1.0f64;
            let mut y = 0.0f64;

            while -2.0 * y < x * x {
                let x_: f64 = rng.sample(Open01);
                let y_: f64 = rng.sample(Open01);

                x = x_.ln() / ziggurat_tables::ZIG_NORM_R;
                y = y_.ln();
            }

            if u < 0.0 {
                x - ziggurat_tables::ZIG_NORM_R
            } else {
                ziggurat_tables::ZIG_NORM_R - x
            }
        }

        ziggurat(
            rng,
            true, // this is symmetric
            &ziggurat_tables::ZIG_NORM_X,
            &ziggurat_tables::ZIG_NORM_F,
            pdf,
            zero_case,
        )
    }
}

/// The normal distribution `N(mean, std_dev**2)`.
///
/// This uses the ZIGNOR variant of the Ziggurat method, see [`StandardNormal`]
/// for more details.
///
/// Note that [`StandardNormal`] is an optimised implementation for mean 0, and
/// standard deviation 1.
///
/// # Example
///
/// ```
/// use rand_distr::{Normal, Distribution};
///
/// // mean 2, standard deviation 3
/// let normal = Normal::new(2.0, 3.0).unwrap();
/// let v = normal.sample(&mut rand::thread_rng());
/// println!("{} is from a N(2, 9) distribution", v)
/// ```
///
/// [`StandardNormal`]: crate::StandardNormal
#[derive(Clone, Copy, Debug)]
pub struct Normal<F>
where F: Float, StandardNormal: Distribution<F>
{
    mean: F,
    std_dev: F,
}

/// Error type returned from `Normal::new` and `LogNormal::new`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    /// The mean value is too small (log-normal samples must be positive)
    MeanTooSmall,
    /// The standard deviation or other dispersion parameter is negative,
    /// infinite, or not-a-number.
    BadVariance,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            Error::MeanTooSmall => "mean < 0 or NaN in log-normal distribution",
            Error::BadVariance => "variation parameter is negative or non-finite in (log)normal distribution",
        })
    }
}

#[cfg(feature = "std")]
impl std::error::Error for Error {}

impl<F> Normal<F>
where F: Float, StandardNormal: Distribution<F>
{
    /// Construct, from mean and standard deviation
    ///
    /// Parameters:
    ///
    /// -   mean (`μ`, unrestricted)
    /// -   standard deviation (`σ`, must be finite)
    #[inline]
    pub fn new(mean: F, std_dev: F) -> Result<Normal<F>, Error> {
        if !std_dev.is_finite() {
            return Err(Error::BadVariance);
        }
        Ok(Normal { mean, std_dev })
    }

    /// Construct, from mean and coefficient of variation
    ///
    /// Parameters:
    ///
    /// -   mean (`μ`, unrestricted)
    /// -   coefficient of variation (`cv = abs(σ / μ)`)
    #[inline]
    pub fn from_mean_cv(mean: F, cv: F) -> Result<Normal<F>, Error> {
        if !cv.is_finite() || cv < F::zero() {
            return Err(Error::BadVariance);
        }
        let std_dev = cv * mean;
        Ok(Normal { mean, std_dev })
    }

    /// Sample from a z-score
    ///
    /// This may be useful for generating correlated samples `x1` and `x2`
    /// from two different distributions, as follows.
    /// ```
    /// # use rand::prelude::*;
    /// # use rand_distr::{Normal, StandardNormal};
    /// let mut rng = thread_rng();
    /// let z = StandardNormal.sample(&mut rng);
    /// let x1 = Normal::new(0.0, 1.0).unwrap().from_zscore(z);
    /// let x2 = Normal::new(2.0, -3.0).unwrap().from_zscore(z);
    /// ```
    #[inline]
    pub fn from_zscore(&self, zscore: F) -> F {
        self.mean + self.std_dev * zscore
    }
}

impl<F> Distribution<F> for Normal<F>
where F: Float, StandardNormal: Distribution<F>
{
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
        self.from_zscore(rng.sample(StandardNormal))
    }
}

/// The log-normal distribution `ln N(mean, std_dev**2)`.
///
/// If `X` is log-normal distributed, then `ln(X)` is `N(mean, std_dev**2)`
/// distributed.
///
/// # Example
///
/// ```
/// use rand_distr::{LogNormal, Distribution};
///
/// // mean 2, standard deviation 3
/// let log_normal = LogNormal::new(2.0, 3.0).unwrap();
/// let v = log_normal.sample(&mut rand::thread_rng());
/// println!("{} is from an ln N(2, 9) distribution", v)
/// ```
#[derive(Clone, Copy, Debug)]
pub struct LogNormal<F>
where F: Float, StandardNormal: Distribution<F>
{
    norm: Normal<F>,
}

impl<F> LogNormal<F>
where F: Float, StandardNormal: Distribution<F>
{
    /// Construct, from (log-space) mean and standard deviation
    ///
    /// Parameters are the "standard" log-space measures (these are the mean
    /// and standard deviation of the logarithm of samples):
    ///
    /// -   `mu` (`μ`, unrestricted) is the mean of the underlying distribution
    /// -   `sigma` (`σ`, must be finite) is the standard deviation of the
    ///     underlying Normal distribution
    #[inline]
    pub fn new(mu: F, sigma: F) -> Result<LogNormal<F>, Error> {
        let norm = Normal::new(mu, sigma)?;
        Ok(LogNormal { norm })
    }

    /// Construct, from (linear-space) mean and coefficient of variation
    ///
    /// Parameters are linear-space measures:
    ///
    /// -   mean (`μ > 0`) is the (real) mean of the distribution
    /// -   coefficient of variation (`cv = σ / μ`, requiring `cv ≥ 0`) is a
    ///     standardized measure of dispersion
    ///
    /// As a special exception, `μ = 0, cv = 0` is allowed (samples are `-inf`).
    #[inline]
    pub fn from_mean_cv(mean: F, cv: F) -> Result<LogNormal<F>, Error> {
        if cv == F::zero() {
            let mu = mean.ln();
            let norm = Normal::new(mu, F::zero()).unwrap();
            return Ok(LogNormal { norm });
        }
        if !(mean > F::zero()) {
            return Err(Error::MeanTooSmall);
        }
        if !(cv >= F::zero()) {
            return Err(Error::BadVariance);
        }

        // Using X ~ lognormal(μ, σ), CV² = Var(X) / E(X)²
        // E(X) = exp(μ + σ² / 2) = exp(μ) × exp(σ² / 2)
        // Var(X) = exp(2μ + σ²)(exp(σ²) - 1) = E(X)² × (exp(σ²) - 1)
        // but Var(X) = (CV × E(X))² so CV² = exp(σ²) - 1
        // thus σ² = log(CV² + 1)
        // and exp(μ) = E(X) / exp(σ² / 2) = E(X) / sqrt(CV² + 1)
        let a = F::one() + cv * cv; // e
        let mu = F::from(0.5).unwrap() * (mean * mean / a).ln();
        let sigma = a.ln().sqrt();
        let norm = Normal::new(mu, sigma)?;
        Ok(LogNormal { norm })
    }

    /// Sample from a z-score
    ///
    /// This may be useful for generating correlated samples, as follows.
    /// ```
    /// # use rand::prelude::*;
    /// # use rand_distr::{LogNormal, StandardNormal};
    /// let mut rng = thread_rng();
    /// let z = StandardNormal.sample(&mut rng);
    /// let x1 = LogNormal::from_mean_cv(3.0, 1.0).unwrap().from_zscore(z);
    /// let x2 = LogNormal::from_mean_cv(2.0, 1.0).unwrap().from_zscore(z);
    /// ```
    #[inline]
    pub fn from_zscore(&self, zscore: F) -> F {
        self.norm.from_zscore(zscore).exp()
    }
}

impl<F> Distribution<F> for LogNormal<F>
where F: Float, StandardNormal: Distribution<F>
{
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F {
        self.norm.sample(rng).exp()
    }
}

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

    #[test]
    fn test_normal() {
        let norm = Normal::new(10.0, 10.0).unwrap();
        let mut rng = crate::test::rng(210);
        for _ in 0..1000 {
            norm.sample(&mut rng);
        }
    }
    #[test]
    fn test_normal_cv() {
        let norm = Normal::from_mean_cv(1024.0, 1.0 / 256.0).unwrap();
        assert_eq!((norm.mean, norm.std_dev), (1024.0, 4.0));
    }
    #[test]
    fn test_normal_invalid_sd() {
        assert!(Normal::from_mean_cv(10.0, -1.0).is_err());
    }

    #[test]
    fn test_log_normal() {
        let lnorm = LogNormal::new(10.0, 10.0).unwrap();
        let mut rng = crate::test::rng(211);
        for _ in 0..1000 {
            lnorm.sample(&mut rng);
        }
    }
    #[test]
    fn test_log_normal_cv() {
        let lnorm = LogNormal::from_mean_cv(0.0, 0.0).unwrap();
        assert_eq!((lnorm.norm.mean, lnorm.norm.std_dev), (-std::f64::INFINITY, 0.0));

        let lnorm = LogNormal::from_mean_cv(1.0, 0.0).unwrap();
        assert_eq!((lnorm.norm.mean, lnorm.norm.std_dev), (0.0, 0.0));

        let e = std::f64::consts::E;
        let lnorm = LogNormal::from_mean_cv(e.sqrt(), (e - 1.0).sqrt()).unwrap();
        assert_almost_eq!(lnorm.norm.mean, 0.0, 2e-16);
        assert_almost_eq!(lnorm.norm.std_dev, 1.0, 2e-16);

        let lnorm = LogNormal::from_mean_cv(e.powf(1.5), (e - 1.0).sqrt()).unwrap();
        assert_eq!((lnorm.norm.mean, lnorm.norm.std_dev), (1.0, 1.0));
    }
    #[test]
    fn test_log_normal_invalid_sd() {
        assert!(LogNormal::from_mean_cv(-1.0, 1.0).is_err());
        assert!(LogNormal::from_mean_cv(0.0, 1.0).is_err());
        assert!(LogNormal::from_mean_cv(1.0, -1.0).is_err());
    }
}
<|endoftext|>


<commit_before>//
//
//

use memory::virt::ProtectionMode;
use arch::memory::PAddr;

const KERNEL_TEMP_BASE : usize = 0xFF800000;
const KERNEL_TEMP_LIMIT: usize = 0xFFFF0000;
const KERNEL_BASE_TABLE: usize = 0xFFFF8000;
const KERNEL_TEMP_TABLE: usize = 0xFFFF9000;

pub fn is_fixed_alloc<T>(addr: *const T, size: usize) -> bool {
	const BASE : usize = super::addresses::KERNEL_BASE;
	const LIMIT: usize = super::addresses::KERNEL_BASE + 4*1024*1024;
	let addr = addr as usize;
	if addr < BASE {
		false
	}
	else if addr >= LIMIT {
		false
	}
	else if addr + size > LIMIT {
		false
	}
	else {
		false
		//true
	}
}
// UNSAFE: Can cause aliasing
pub unsafe fn fixed_alloc(p: PAddr, count: usize) -> Option<*mut ()> {
	None
}

#[derive(Copy,Clone)]
enum PageEntryRegion {
	NonGlobal,
	Global,
}
impl PageEntryRegion {
	unsafe fn get_section_ent(&self, idx: usize) -> &mut u32 {
		assert!(idx < 4096);
		match self
		{
		&PageEntryRegion::NonGlobal => todo!("PageEntryRegion::get_section_ent - non-global"),
		&PageEntryRegion::Global => &mut *((KERNEL_BASE_TABLE + idx * 4) as *mut u32),
		}
	}
}
enum PageEntry {
	Section {
		rgn: PageEntryRegion,
		idx: usize,
		ofs: usize
		},
	Page {
		mapping: TempHandle,
		idx: usize,
		ofs: usize
		},
}
impl PageEntry
{
	fn alloc(addr: *const (), level: usize) -> Result<PageEntry, ()> {
		todo!("PageEntry::alloc");
	}
	fn get(addr: *const ()) -> PageEntry {
		use super::addresses::KERNEL_BASE;
		let (rgn, p_idx) = if (addr as usize) < KERNEL_BASE {
				(PageEntryRegion::NonGlobal, (addr as usize - KERNEL_BASE) >> 12)
			}
			else {
				(PageEntryRegion::Global, (addr as usize) >> 12)
			};

		// SAFE: Aliasing in this case is benign
		let sect_ent = unsafe { *rgn.get_section_ent(p_idx >> 8) };
		if sect_ent & 0b11 == 0b01 {
			PageEntry::Page {
				// SAFE: ... won't be mutated (yet), but need to ensure safety TODO
				mapping: unsafe { TempHandle::new( sect_ent & !0xFFF ) },
				idx: p_idx,
				ofs: (addr as usize) & 0xFFF,
				}
		}
		else {
			PageEntry::Section {
				rgn: rgn,
				idx: p_idx >> 8,
				ofs: (addr as usize) & 0xFF_FFF,
				}
		}
	}

	fn is_reserved(&self) -> bool {
		// TODO: Need to _ensure_ that the page table is not removed during manipulation
		// SAFE: Aliasing is benign, and page table should be mapped (see above TODO)
		unsafe {
			match self
			{
			&PageEntry::Section { rgn, idx, .. } => (*rgn.get_section_ent(idx) & 3 != 0),
			&PageEntry::Page { ref mapping, idx, .. } => (mapping[idx & 0x3FF] & 3 != 0),
			}
		}
	}

	fn phys_addr(&self) -> ::arch::memory::PAddr {
		// SAFE: Aliasing is benign, and page table should be mapped (see above TODO)
		unsafe {
			match self
			{
			&PageEntry::Section { rgn, idx, .. } => (*rgn.get_section_ent(idx) & !0xFFF),
			&PageEntry::Page { ref mapping, idx, .. } => (mapping[idx & 0x3FF] & !0xFFF),
			}
		}
	}
}

struct AtomicU32(::core::cell::UnsafeCell<u32>);
impl AtomicU32 {
	pub fn cxchg(&self, val: u32, new: u32) -> u32 {
		// SAFE: Atomic
		unsafe { ::core::intrinsics::atomic_cxchg_relaxed(self.0.get(), val, new) }
	}
	pub fn store(&self, val: u32) {
		// SAFE: Atomic
		unsafe { ::core::intrinsics::atomic_store_relaxed(self.0.get(), val) }
	}
	pub fn load(&self) -> u32 {
		// SAFE: Atomic
		unsafe { ::core::intrinsics::atomic_load_relaxed(self.0.get()) }
	}
}
extern "C" {
	static kernel_table0: [AtomicU32; 0x800*2];
	static kernel_exception_map: [AtomicU32; 1024];
}

fn get_table_addr<T>(a: *const T, alloc: bool) -> Option< (::arch::memory::PAddr, usize) > {
	let addr = a as usize;
	let page = addr >> 12;
	let (ttbr_ofs, tab_idx) = (page >> 8, page & 0xFF);
	let ent_r = if ttbr_ofs < 0x800 {
			todo!("get_table_addr - User");
		}
		else {
			// Kernel
			&kernel_table0[ ttbr_ofs ]
		};

	//let ent_v = ent_r.load();
	let ent_v = ent_r.load();
	match ent_v & 0xFFF
	{
	0 => if alloc {
			let frame = ::memory::phys::allocate_bare().expect("TODO get_table_addr");
			let ent_v = ent_r.cxchg(0, frame + 0x1);
			if ent_v != 0 {
				::memory::phys::deref_frame(frame);
				Some( (ent_v & !0xFFF, tab_idx) )
			}
			else {
				Some( (frame & !0xFFF, tab_idx) )
			}
		}
		else {
			None
		},
	1 => Some( (ent_v & !0xFFF, tab_idx) ),
	v @ _ => todo!("get_table_addr - Other flags bits {:#x}", v),
	}
}
//static S_TEMP_MAP_SEMAPHORE: Semaphore = Semaphore::new();
const TEMP_BASE: usize = 0xFFC0_0000;
struct TempHandle(*mut [u32; 1024]);
impl TempHandle
{
	/// UNSAFE: User must ensure that address is valid, and that no aliasing occurs
	unsafe fn new(phys: ::arch::memory::PAddr) -> TempHandle {
		let val = (phys as u32) + 0x13;

		//S_TEMP_MAP_SEMAPHORE.take();
		// #1023 is reserved for -1 mapping
		for i in 0 .. 1023 {
			if kernel_exception_map[i].cxchg(0, val) == 0 {
				return TempHandle( (TEMP_BASE + i * 0x1000) as *mut _ );
			}
		}
		panic!("No free temp mappings");
	}
}
impl ::core::ops::Deref for TempHandle {
	type Target = [u32];
	fn deref(&self) -> &[u32] {
		// SAFE: We should have unique access
		unsafe { &*self.0 }
	}
}
impl ::core::ops::DerefMut for TempHandle {
	fn deref_mut(&mut self) -> &mut [u32] {
		// SAFE: We should have unique access
		unsafe { &mut *self.0 }
	}
}
impl ::core::ops::Drop for TempHandle {
	fn drop(&mut self) {
		let i = (self.0 as usize - TEMP_BASE) / 0x1000;
		kernel_exception_map[i].store(0);
		//S_TEMP_MAP_SEMAPHORE.add();
	}
}

pub fn is_reserved<T>(addr: *const T) -> bool {
	PageEntry::get(addr as *const ()).is_reserved()
}
pub fn get_phys<T>(addr: *const T) -> ::arch::memory::PAddr {
	PageEntry::get(addr as *const ()).phys_addr()
}

pub fn get_info<T>(addr: *const T) -> Option<(u32, ::memory::virt::ProtectionMode)> {
	todo!("get_info")
}

pub unsafe fn map(a: *mut (), p: PAddr, mode: ProtectionMode) {
	// 1. Map the relevant table in the temp area
	let (tab_phys, idx) = get_table_addr(a, true).unwrap();
	// TODO: Ensure nothing else is manipulating this segment of AS
	let mut mh = TempHandle::new( tab_phys );
	// 2. Insert
	let mode_flags = match mode
		{
		ProtectionMode::Unmapped => panic!("Invalid pass of Unmapped to map"),
		ProtectionMode::KernelRO => 0x212,
		ProtectionMode::KernelRW => 0x012,
		ProtectionMode::KernelRX => 0x053,
		ProtectionMode::UserRO => 0x232,
		ProtectionMode::UserRW => 0x032,
		ProtectionMode::UserRX => 0x233,
		ProtectionMode::UserRWX => 0x033,
		ProtectionMode::UserCOW => 0x223,	// 1,10 is a deprecated encoding, need to find a better encoding
		};
	mh[idx] = p + mode_flags;
}
pub unsafe fn reprotect(a: *mut (), mode: ProtectionMode) {
	todo!("reprotect({:p}, {:?}", a, mode)
}
pub unsafe fn unmap(a: *mut ()) -> Option<PAddr> {
	todo!("unmap")
}

#[derive(Debug)]
pub struct AddressSpace(u32);
impl AddressSpace
{
	pub fn pid0() -> AddressSpace {
		extern "C" {
			static kernel_table0_ofs: ::Void;
			static kernel_data_start: u32;
		}
		let tab0_addr = kernel_data_start + (&kernel_table0_ofs as *const _ as usize as u32);
		AddressSpace( tab0_addr )
	}
	pub fn new(clone_start: usize, clone_end: usize) -> Result<AddressSpace,::memory::virt::MapError> {
		todo!("AddressSpace::new({:#x} -- {:#x})", clone_start, clone_end);
	}

	pub fn get_ttbr0(&self) -> u32 { self.0 }
}

<commit_msg>Kernel/ARMv7 - Minor cleanup to VMM code<commit_after>//
//
//

use memory::virt::ProtectionMode;
use arch::memory::PAddr;

const KERNEL_BASE_TABLE: usize = 0xFFFF8000;

pub fn is_fixed_alloc<T>(_addr: *const T, _size: usize) -> bool {
	//const BASE : usize = super::addresses::KERNEL_BASE;
	//const ONEMEG: usize =1024*1024
	//const LIMIT: usize = super::addresses::KERNEL_BASE + 4*ONEMEG;
	false
}
// UNSAFE: Can cause aliasing
pub unsafe fn fixed_alloc(_p: PAddr, _count: usize) -> Option<*mut ()> {
	None
}

#[derive(Copy,Clone)]
enum PageEntryRegion {
	NonGlobal,
	Global,
}
impl PageEntryRegion {
	unsafe fn get_section_ent(&self, idx: usize) -> &mut u32 {
		assert!(idx < 4096);
		match self
		{
		&PageEntryRegion::NonGlobal => todo!("PageEntryRegion::get_section_ent - non-global"),
		&PageEntryRegion::Global => &mut *((KERNEL_BASE_TABLE + idx * 4) as *mut u32),
		}
	}
}
enum PageEntry {
	Section {
		rgn: PageEntryRegion,
		idx: usize,
		ofs: usize
		},
	Page {
		mapping: TempHandle,
		idx: usize,
		ofs: usize
		},
}
impl PageEntry
{
	fn alloc(addr: *const (), level: usize) -> Result<PageEntry, ()> {
		todo!("PageEntry::alloc({:p}, level={})", addr, level);
	}
	/// Obtain a page entry for the specified address
	fn get(addr: *const ()) -> PageEntry {
		use super::addresses::KERNEL_BASE;
		let (rgn, p_idx) = if (addr as usize) < KERNEL_BASE {
				(PageEntryRegion::NonGlobal, (addr as usize - KERNEL_BASE) >> 12)
			}
			else {
				(PageEntryRegion::Global, (addr as usize) >> 12)
			};

		// SAFE: Aliasing in this case is benign
		let sect_ent = unsafe { *rgn.get_section_ent(p_idx >> 8) };
		if sect_ent & 0b11 == 0b01 {
			PageEntry::Page {
				// SAFE: ... won't be mutated (yet), but need to ensure safety TODO
				mapping: unsafe { TempHandle::new( sect_ent & !0xFFF ) },
				idx: p_idx,
				ofs: (addr as usize) & 0xFFF,
				}
		}
		else {
			PageEntry::Section {
				rgn: rgn,
				idx: p_idx >> 8,
				ofs: (addr as usize) & 0xFF_FFF,
				}
		}
	}

	fn is_reserved(&self) -> bool {
		// TODO: Need to _ensure_ that the page table is not removed during manipulation
		// SAFE: Aliasing is benign, and page table should be mapped (see above TODO)
		unsafe {
			match self
			{
			&PageEntry::Section { rgn, idx, .. } => (*rgn.get_section_ent(idx) & 3 != 0),
			&PageEntry::Page { ref mapping, idx, .. } => (mapping[idx & 0x3FF] & 3 != 0),
			}
		}
	}

	fn phys_addr(&self) -> ::arch::memory::PAddr {
		// SAFE: Aliasing is benign, and page table should be mapped (see above TODO)
		unsafe {
			match self
			{
			&PageEntry::Section { rgn, idx, ofs } => (*rgn.get_section_ent(idx) & !0xFFF) + ofs as u32,
			&PageEntry::Page { ref mapping, idx ,ofs } => (mapping[idx & 0x3FF] & !0xFFF) + ofs as u32,
			}
		}
	}
}

#[repr(C)]
struct AtomicU32(::core::cell::UnsafeCell<u32>);
impl AtomicU32 {
	pub fn cxchg(&self, val: u32, new: u32) -> u32 {
		// SAFE: Atomic
		unsafe { ::core::intrinsics::atomic_cxchg_relaxed(self.0.get(), val, new) }
	}
	pub fn store(&self, val: u32) {
		// SAFE: Atomic
		unsafe { ::core::intrinsics::atomic_store_relaxed(self.0.get(), val) }
	}
	pub fn load(&self) -> u32 {
		// SAFE: Atomic
		unsafe { ::core::intrinsics::atomic_load_relaxed(self.0.get()) }
	}
}
extern "C" {
	static kernel_table0: [AtomicU32; 0x800*2];
	static kernel_exception_map: [AtomicU32; 1024];
}

fn get_table_addr<T>(a: *const T, alloc: bool) -> Option< (::arch::memory::PAddr, usize) > {
	let addr = a as usize;
	let page = addr >> 12;
	let (ttbr_ofs, tab_idx) = (page >> 8, page & 0xFF);
	let ent_r = if ttbr_ofs < 0x800 {
			todo!("get_table_addr - User");
		}
		else {
			// Kernel
			&kernel_table0[ ttbr_ofs ]
		};

	//let ent_v = ent_r.load();
	let ent_v = ent_r.load();
	match ent_v & 0xFFF
	{
	0 => if alloc {
			let frame = ::memory::phys::allocate_bare().expect("TODO get_table_addr");
			let ent_v = ent_r.cxchg(0, frame + 0x1);
			if ent_v != 0 {
				::memory::phys::deref_frame(frame);
				Some( (ent_v & !0xFFF, tab_idx) )
			}
			else {
				Some( (frame & !0xFFF, tab_idx) )
			}
		}
		else {
			None
		},
	1 => Some( (ent_v & !0xFFF, tab_idx) ),
	v @ _ => todo!("get_table_addr - Other flags bits {:#x}", v),
	}
}
//static S_TEMP_MAP_SEMAPHORE: Semaphore = Semaphore::new();
const KERNEL_TEMP_BASE : usize = 0xFFC00000;

struct TempHandle(*mut [u32; 1024]);
impl TempHandle
{
	/// UNSAFE: User must ensure that address is valid, and that no aliasing occurs
	unsafe fn new(phys: ::arch::memory::PAddr) -> TempHandle {
		let val = (phys as u32) + 0x13;

		//S_TEMP_MAP_SEMAPHORE.take();
		// #1023 is reserved for -1 mapping
		for i in 0 .. 1023 {
			if kernel_exception_map[i].cxchg(0, val) == 0 {
				return TempHandle( (KERNEL_TEMP_BASE + i * 0x1000) as *mut _ );
			}
		}
		panic!("No free temp mappings");
	}
}
impl ::core::ops::Deref for TempHandle {
	type Target = [u32];
	fn deref(&self) -> &[u32] {
		// SAFE: We should have unique access
		unsafe { &*self.0 }
	}
}
impl ::core::ops::DerefMut for TempHandle {
	fn deref_mut(&mut self) -> &mut [u32] {
		// SAFE: We should have unique access
		unsafe { &mut *self.0 }
	}
}
impl ::core::ops::Drop for TempHandle {
	fn drop(&mut self) {
		let i = (self.0 as usize - KERNEL_TEMP_BASE) / 0x1000;
		kernel_exception_map[i].store(0);
		//S_TEMP_MAP_SEMAPHORE.add();
	}
}

pub fn is_reserved<T>(addr: *const T) -> bool {
	PageEntry::get(addr as *const ()).is_reserved()
}
pub fn get_phys<T>(addr: *const T) -> ::arch::memory::PAddr {
	PageEntry::get(addr as *const ()).phys_addr()
}

pub fn get_info<T>(addr: *const T) -> Option<(u32, ::memory::virt::ProtectionMode)> {
	todo!("get_info")
}

pub unsafe fn map(a: *mut (), p: PAddr, mode: ProtectionMode) {
	// 1. Map the relevant table in the temp area
	let (tab_phys, idx) = get_table_addr(a, true).unwrap();
	// TODO: Ensure nothing else is manipulating this segment of AS
	let mut mh = TempHandle::new( tab_phys );
	// 2. Insert
	let mode_flags = match mode
		{
		ProtectionMode::Unmapped => panic!("Invalid pass of Unmapped to map"),
		ProtectionMode::KernelRO => 0x212,
		ProtectionMode::KernelRW => 0x012,
		ProtectionMode::KernelRX => 0x053,
		ProtectionMode::UserRO => 0x232,
		ProtectionMode::UserRW => 0x032,
		ProtectionMode::UserRX => 0x233,
		ProtectionMode::UserRWX => 0x033,
		ProtectionMode::UserCOW => 0x223,	// 1,10 is a deprecated encoding for RO, need to find a better encoding
		};
	mh[idx] = p + mode_flags;
}
pub unsafe fn reprotect(a: *mut (), mode: ProtectionMode) {
	todo!("reprotect({:p}, {:?}", a, mode)
}
pub unsafe fn unmap(a: *mut ()) -> Option<PAddr> {
	todo!("unmap")
}

#[derive(Debug)]
pub struct AddressSpace(u32);
impl AddressSpace
{
	pub fn pid0() -> AddressSpace {
		extern "C" {
			static kernel_table0_ofs: ::Void;
			static kernel_data_start: u32;
		}
		let tab0_addr = kernel_data_start + (&kernel_table0_ofs as *const _ as usize as u32);
		AddressSpace( tab0_addr )
	}
	pub fn new(clone_start: usize, clone_end: usize) -> Result<AddressSpace,::memory::virt::MapError> {
		todo!("AddressSpace::new({:#x} -- {:#x})", clone_start, clone_end);
	}

	pub fn get_ttbr0(&self) -> u32 { self.0 }
}

<|endoftext|>


<commit_before>use std::collections::hashmap::HashMap;
use time;
use std::mem;
use sync::comm::Receiver;
use image::{GenericImage, SubImage, Pixel};
use image::{ImageBuf, Rgba};
use image;
use glfw;
use glfw::Context as _Context;
use gfx;
use gfx::tex;
use gfx::{Device, DeviceHelper, ToSlice, CommandBuffer};
use gfx::{GlDevice};
use key;
use atlas::{AtlasBuilder, Atlas};
use util;
use color;
use color::{Rgb};
use geom::{V2};

static FONT_DATA: &'static [u8] = include_bin!("../assets/font.png");

pub struct Canvas {
    title: String,
    dim: V2<u32>,
    frame_interval: Option<f64>,
    builder: AtlasBuilder,
    font_glyphs: HashMap<char, Image>,
}

/// Toplevel graphics drawing and input reading context.
impl Canvas {
    pub fn new() -> Canvas {
        let mut ret = Canvas {
            title: "window".to_string(),
            dim: V2(640, 360),
            frame_interval: None,
            builder: AtlasBuilder::new(),
            font_glyphs: HashMap::new(),
        };
        ret.init_font();
        ret
    }

    pub fn set_title(mut self, title: &str) -> Canvas {
        self.title = title.to_string();
        self
    }

    /// Set the frame rate.
    pub fn set_frame_interval(mut self, interval_s: f64) -> Canvas {
        assert!(interval_s > 0.00001);
        self.frame_interval = Some(interval_s);
        self
    }

    /// Set the resolution.
    pub fn set_dim(mut self, dim: V2<u32>) -> Canvas {
        self.dim = dim;
        self
    }

    pub fn add_image<P: Pixel<u8>, I: GenericImage<P>>(
        &mut self, image: I) -> Image {
        Image(self.builder.push(image))
    }

    /// Start running the engine, return an event iteration.
    pub fn run(&mut self) -> Context {
        Context::new(
            self.dim,
            self.title.as_slice(),
            Some(1.0 / 30.0),
            Atlas::new(&self.builder))
    }

    fn init_font(&mut self) {
        let mut font_sheet = util::color_key(
            &image::load_from_memory(FONT_DATA, image::PNG).unwrap(),
            0x80u8, 0x80u8, 0x80u8);
        for i in range(0u32, 96u32) {
            let x = 8u32 * (i % 16u32);
            let y = 8u32 * (i / 16u32);
            let glyph = self.add_image(SubImage::new(&mut font_sheet, x, y, 8, 8));
            self.font_glyphs.insert((i + 32) as u8 as char, glyph);
        }
    }
}

/// Interface to render to a live display.
pub struct Context {
    glfw: glfw::Glfw,
    window: glfw::Window,
    events: Receiver<(f64, glfw::WindowEvent)>,
    graphics: gfx::Graphics<gfx::GlDevice, gfx::GlCommandBuffer>,
    frame: gfx::Frame,

    state: State,
    frame_interval: Option<f64>,
    last_render_time: f64,
    atlas: Atlas,
    atlas_tex: Texture,
    flatshade: gfx::TextureHandle,
    resolution: V2<u32>,
}

#[deriving(PartialEq)]
enum State {
    Normal,
    EndFrame,
}

impl Context {
    fn new(
        dim: V2<u32>,
        title: &str,
        frame_interval: Option<f64>,
        atlas: Atlas) -> Context {

        let glfw = glfw::init(glfw::FAIL_ON_ERRORS).unwrap();
        let (window, events) = glfw
            .create_window(dim.0, dim.1,
            title.as_slice(), glfw::Windowed)
            .expect("Failed to open window");
        window.make_current();
        glfw.set_error_callback(glfw::FAIL_ON_ERRORS);
        window.set_key_polling(true);
        window.set_char_polling(true);

        let device = gfx::GlDevice::new(|s| window.get_proc_address(s));
        let mut graphics = gfx::Graphics::new(device);
        //let frame = gfx::Frame::new(dim.0 as u16, dim.1 as u16);
        let (w, h) = window.get_framebuffer_size();
        let frame = gfx::Frame::new(w as u16, h as u16);
        let atlas_tex = Texture::from_rgba8(&atlas.image, &mut graphics.device);

        // Blank white texture so we can draw flatshaded stuff without
        // switching to non-texturing shader.
        let mut inf = tex::TextureInfo::new();
        inf.width = 1;
        inf.height = 1;
        inf.kind = tex::Texture2D;
        inf.format = tex::RGBA8;
        let flatshade = graphics.device.create_texture(inf).unwrap();
        graphics.device.update_texture(&flatshade, &inf.to_image_info(), &[0xffu8, 0xff, 0xff, 0xff]).unwrap();

        Context {
            glfw: glfw,
            window: window,
            events: events,
            graphics: graphics,
            frame: frame,
            state: Normal,
            frame_interval: frame_interval,
            last_render_time: time::precise_time_s(),
            atlas: atlas,
            atlas_tex: atlas_tex,
            flatshade: flatshade,
            resolution: dim,
        }
    }

    /// Clear the screen
    pub fn clear(&mut self, color: &Rgb) {
        self.graphics.clear(
            gfx::ClearData {
                color: color.to_array(),
                depth: 1.0,
                stencil: 0,
            }, gfx::Color, &self.frame);
    }

    /// Mess with drawy stuff
    pub fn draw_test(&mut self) {
        let img = self.font_image('X').unwrap();
        self.draw_image(V2(1, 1), img, &color::ALICEBLUE);
        self.draw_line(V2(10, 10), V2(100, 50), 3f32, &color::YELLOW);
    }

    pub fn draw_image(&mut self, offset: V2<int>, Image(idx): Image, color: &Rgb) {
        let scale = self.resolution.map(|x| 2.0 / (x as f32));

        let texcoords = self.atlas.texcoords[idx];
        let V2(u1, v1) = texcoords.min();
        let V2(u2, v2) = texcoords.max();

        let vertices = self.atlas.vertices[idx];
        let V2(x1, y1) = (vertices.min() + offset.map(|x| x as f32)).mul(scale) - V2(1f32, 1f32);
        let V2(x2, y2) = (vertices.max() + offset.map(|x| x as f32)).mul(scale) - V2(1f32, 1f32);

        let mesh = self.graphics.device.create_mesh([
            Vertex { pos: [x1, -y2], tex_coord: [u1, v2] },
            Vertex { pos: [x2, -y2], tex_coord: [u2, v2] },
            Vertex { pos: [x1, -y1], tex_coord: [u1, v1] },

            Vertex { pos: [x2, -y2], tex_coord: [u2, v2] },
            Vertex { pos: [x1, -y1], tex_coord: [u1, v1] },
            Vertex { pos: [x2, -y1], tex_coord: [u2, v1] },
        ]);

        let sampler_info = Some(self.graphics.device.create_sampler(
            tex::SamplerInfo::new(tex::Scale, tex::Clamp)));
        let params = ShaderParam {
            color: color.to_array(),
            s_texture: (self.atlas_tex.tex, sampler_info),
        };

        let slice = mesh.to_slice(gfx::TriangleList);
        let program = self.graphics.device.link_program(
            VERTEX_SRC.clone(), FRAGMENT_SRC.clone()).unwrap();
        let batch: gfx::batch::RefBatch<_ShaderParamLink, ShaderParam> = self.graphics.make_batch(
            &program, &mesh, slice, &gfx::DrawState::new()).unwrap();
        self.graphics.draw(&batch, &params, &self.frame);
    }

    pub fn draw_line(&mut self, p1: V2<int>, p2: V2<int>, thickness: f32, color: &Rgb) {
        let scale = self.resolution.map(|x| 2.0 / (x as f32));
        let p1 = p1.map(|x| x as f32).mul(scale) - V2(1f32, 1f32);
        let p2 = p2.map(|x| x as f32).mul(scale) - V2(1f32, 1f32);

        let mesh = self.graphics.device.create_mesh([
            Vertex { pos: [p1.0, -p1.1], tex_coord: [0.0, 0.0] },
            Vertex { pos: [p2.0, -p2.1], tex_coord: [0.0, 0.0] },
        ]);

        let sampler_info = Some(self.graphics.device.create_sampler(
            tex::SamplerInfo::new(tex::Scale, tex::Clamp)));
        let params = ShaderParam {
            color: color.to_array(),
            s_texture: (self.flatshade, sampler_info),
        };
        let slice = mesh.to_slice(gfx::Line);
        let program = self.graphics.device.link_program(
            VERTEX_SRC.clone(), FRAGMENT_SRC.clone()).unwrap();
        let mut draw_state = gfx::DrawState::new();
        draw_state.primitive.method = gfx::state::Line(thickness);
        let batch: gfx::batch::RefBatch<_ShaderParamLink, ShaderParam> = self.graphics.make_batch(
            &program, &mesh, slice, &draw_state).unwrap();
        self.graphics.draw(&batch, &params, &self.frame);
    }

    pub fn font_image(&self, c: char) -> Option<Image> {
        let idx = c as int;
        if idx >= 32 && idx < 128 {
            Some(Image((idx - 32) as uint))
        } else {
            None
        }
    }

    pub fn image_dim(&self, Image(idx): Image) -> V2<uint> {
        self.atlas.vertices[idx].1.map(|x| x as uint)
    }
}

pub enum Event<'a> {
    /// Time to render the screen. Call your own render code on the Context
    /// value when you get this.
    Render(&'a mut Context),
    /// Some printable text entered. Data is whole strings to account for
    /// exotic input devices.
    Text(String),
    KeyPressed(key::Key),
    KeyReleased(key::Key),
}

impl<'a> Iterator<Event<'a>> for Context {
    fn next(&mut self) -> Option<Event<'a>> {
        // After a render event, control will return here on a new
        // iter call. Do post-render work here.
        if self.state == EndFrame {
            self.state = Normal;
            self.graphics.end_frame();
            self.window.swap_buffers();
        }

        loop {
            if self.window.should_close() {
                return None;
            }

            self.glfw.poll_events();

            match self.events.try_recv() {
                Ok((_, event)) => {
                    match event {
                        glfw::CharEvent(ch) => {
                            return Some(Text(String::from_char(1, ch)));
                        }
                        glfw::KeyEvent(k, _scan, action, _mods) => {
                            match key::translate_glfw_key(k).map(|k| {
                                if action == glfw::Press || action == glfw::Repeat {
                                    KeyPressed(k)
                                }
                                else {
                                    KeyReleased(k)
                                }
                            }) {
                                Some(e) => { return Some(e); }
                                _ => ()
                            }
                        }
                        // TODO Mouse events.
                        _ => ()
                    }
                }
                _ => ()
            }

            let t = time::precise_time_s();
            if self.frame_interval.map_or(true,
                |x| t - self.last_render_time >= x) {
                self.last_render_time = t;

                // Time to render, must return a handle to self.
                // XXX: Need unsafe hackery to get around lifetimes check.
                self.state = EndFrame;
                let (w, h) = self.window.get_framebuffer_size();
                self.frame = gfx::Frame::new(w as u16, h as u16);
                unsafe {
                    return Some(Render(mem::transmute(self)))
                }
            }
        }
    }
}

/// Drawable images stored in the Canvas.
#[deriving(Clone, PartialEq)]
pub struct Image(uint);

static VERTEX_SRC: gfx::ShaderSource = shaders! {
GLSL_120: b"
    #version 120

    uniform vec4 color;

    attribute vec2 a_pos;
    attribute vec2 a_tex_coord;

    // TODO: Make color a uniform argument.
    varying vec2 v_tex_coord;
    varying vec4 v_color;

    void main() {
        v_tex_coord = a_tex_coord;
        v_color = color;
        gl_Position = vec4(a_pos, 0.0, 1.0);
    }
"
};

static FRAGMENT_SRC: gfx::ShaderSource = shaders! {
GLSL_120: b"
    #version 120

    uniform sampler2D s_texture;

    varying vec2 v_tex_coord;
    varying vec4 v_color;

    void main() {
        vec4 tex_color = texture2D(s_texture, v_tex_coord);
        if (tex_color.a == 0.0) discard;
        gl_FragColor = v_color * tex_color;
    }
"
};

#[shader_param(Program)]
pub struct ShaderParam {
    pub color: [f32, ..4],
    pub s_texture: gfx::shade::TextureParam,
}

#[vertex_format]
struct Vertex {
    #[name = "a_pos"]
    pos: [f32, ..2],

    #[name = "a_tex_coord"]
    tex_coord: [f32, ..2],
}

impl Clone for Vertex {
    fn clone(&self) -> Vertex { *self }
}

struct Texture {
    tex: gfx::TextureHandle,
    pub width: u32,
    pub height: u32,
}

impl Texture {
    fn from_rgba8<D: Device<C>, C: CommandBuffer>(
        img: &ImageBuf<Rgba<u8>>,
        d: &mut D) -> Texture {
        let (w, h) = img.dimensions();
        let mut info = tex::TextureInfo::new();
        info.width = w as u16;
        info.height = h as u16;
        info.kind = tex::Texture2D;
        info.format = tex::RGBA8;

        let tex = d.create_texture(info).unwrap();
        d.update_texture(&tex, &info.to_image_info(), img.pixelbuf()).unwrap();

        Texture {
            tex: tex,
            width: w,
            height: h,
        }
    }
}
<commit_msg>Streamline font loading code<commit_after>use std::collections::hashmap::HashMap;
use time;
use std::mem;
use sync::comm::Receiver;
use image::{GenericImage, SubImage, Pixel};
use image::{ImageBuf, Rgba};
use image;
use glfw;
use glfw::Context as _Context;
use gfx;
use gfx::tex;
use gfx::{Device, DeviceHelper, ToSlice, CommandBuffer};
use gfx::{GlDevice};
use key;
use atlas::{AtlasBuilder, Atlas};
use util;
use color;
use color::{Rgb};
use geom::{V2};

pub struct Canvas {
    title: String,
    dim: V2<u32>,
    frame_interval: Option<f64>,
    builder: AtlasBuilder,
    font_glyphs: HashMap<char, Image>,
}

/// Toplevel graphics drawing and input reading context.
impl Canvas {
    pub fn new() -> Canvas {
        let mut ret = Canvas {
            title: "window".to_string(),
            dim: V2(640, 360),
            frame_interval: None,
            builder: AtlasBuilder::new(),
            font_glyphs: HashMap::new(),
        };
        ret.init_font();
        ret
    }

    pub fn set_title(mut self, title: &str) -> Canvas {
        self.title = title.to_string();
        self
    }

    /// Set the frame rate.
    pub fn set_frame_interval(mut self, interval_s: f64) -> Canvas {
        assert!(interval_s > 0.00001);
        self.frame_interval = Some(interval_s);
        self
    }

    /// Set the resolution.
    pub fn set_dim(mut self, dim: V2<u32>) -> Canvas {
        self.dim = dim;
        self
    }

    pub fn add_image<P: Pixel<u8>, I: GenericImage<P>>(
        &mut self, image: I) -> Image {
        Image(self.builder.push(image))
    }

    /// Start running the engine, return an event iteration.
    pub fn run(&mut self) -> Context {
        Context::new(
            self.dim,
            self.title.as_slice(),
            Some(1.0 / 30.0),
            Atlas::new(&self.builder))
    }

    fn init_font(&mut self) {
        let mut font_sheet = util::color_key(
            &image::load_from_memory(include_bin!("../assets/font.png"), image::PNG).unwrap(),
            0x80u8, 0x80u8, 0x80u8);
        for i in range(0u32, 96u32) {
            let x = 8u32 * (i % 16u32);
            let y = 8u32 * (i / 16u32);
            let glyph = self.add_image(SubImage::new(&mut font_sheet, x, y, 8, 8));
            self.font_glyphs.insert((i + 32) as u8 as char, glyph);
        }
    }
}

/// Interface to render to a live display.
pub struct Context {
    glfw: glfw::Glfw,
    window: glfw::Window,
    events: Receiver<(f64, glfw::WindowEvent)>,
    graphics: gfx::Graphics<gfx::GlDevice, gfx::GlCommandBuffer>,
    frame: gfx::Frame,

    state: State,
    frame_interval: Option<f64>,
    last_render_time: f64,
    atlas: Atlas,
    atlas_tex: Texture,
    flatshade: gfx::TextureHandle,
    resolution: V2<u32>,
}

#[deriving(PartialEq)]
enum State {
    Normal,
    EndFrame,
}

impl Context {
    fn new(
        dim: V2<u32>,
        title: &str,
        frame_interval: Option<f64>,
        atlas: Atlas) -> Context {

        let glfw = glfw::init(glfw::FAIL_ON_ERRORS).unwrap();
        let (window, events) = glfw
            .create_window(dim.0, dim.1,
            title.as_slice(), glfw::Windowed)
            .expect("Failed to open window");
        window.make_current();
        glfw.set_error_callback(glfw::FAIL_ON_ERRORS);
        window.set_key_polling(true);
        window.set_char_polling(true);

        let device = gfx::GlDevice::new(|s| window.get_proc_address(s));
        let mut graphics = gfx::Graphics::new(device);
        //let frame = gfx::Frame::new(dim.0 as u16, dim.1 as u16);
        let (w, h) = window.get_framebuffer_size();
        let frame = gfx::Frame::new(w as u16, h as u16);
        let atlas_tex = Texture::from_rgba8(&atlas.image, &mut graphics.device);

        // Blank white texture so we can draw flatshaded stuff without
        // switching to non-texturing shader.
        let mut inf = tex::TextureInfo::new();
        inf.width = 1;
        inf.height = 1;
        inf.kind = tex::Texture2D;
        inf.format = tex::RGBA8;
        let flatshade = graphics.device.create_texture(inf).unwrap();
        graphics.device.update_texture(&flatshade, &inf.to_image_info(), &[0xffu8, 0xff, 0xff, 0xff]).unwrap();

        Context {
            glfw: glfw,
            window: window,
            events: events,
            graphics: graphics,
            frame: frame,
            state: Normal,
            frame_interval: frame_interval,
            last_render_time: time::precise_time_s(),
            atlas: atlas,
            atlas_tex: atlas_tex,
            flatshade: flatshade,
            resolution: dim,
        }
    }

    /// Clear the screen
    pub fn clear(&mut self, color: &Rgb) {
        self.graphics.clear(
            gfx::ClearData {
                color: color.to_array(),
                depth: 1.0,
                stencil: 0,
            }, gfx::Color, &self.frame);
    }

    /// Mess with drawy stuff
    pub fn draw_test(&mut self) {
        let img = self.font_image('X').unwrap();
        self.draw_image(V2(1, 1), img, &color::ALICEBLUE);
        self.draw_line(V2(10, 10), V2(100, 50), 3f32, &color::YELLOW);
    }

    pub fn draw_image(&mut self, offset: V2<int>, Image(idx): Image, color: &Rgb) {
        let scale = self.resolution.map(|x| 2.0 / (x as f32));

        let texcoords = self.atlas.texcoords[idx];
        let V2(u1, v1) = texcoords.min();
        let V2(u2, v2) = texcoords.max();

        let vertices = self.atlas.vertices[idx];
        let V2(x1, y1) = (vertices.min() + offset.map(|x| x as f32)).mul(scale) - V2(1f32, 1f32);
        let V2(x2, y2) = (vertices.max() + offset.map(|x| x as f32)).mul(scale) - V2(1f32, 1f32);

        let mesh = self.graphics.device.create_mesh([
            Vertex { pos: [x1, -y2], tex_coord: [u1, v2] },
            Vertex { pos: [x2, -y2], tex_coord: [u2, v2] },
            Vertex { pos: [x1, -y1], tex_coord: [u1, v1] },

            Vertex { pos: [x2, -y2], tex_coord: [u2, v2] },
            Vertex { pos: [x1, -y1], tex_coord: [u1, v1] },
            Vertex { pos: [x2, -y1], tex_coord: [u2, v1] },
        ]);

        let sampler_info = Some(self.graphics.device.create_sampler(
            tex::SamplerInfo::new(tex::Scale, tex::Clamp)));
        let params = ShaderParam {
            color: color.to_array(),
            s_texture: (self.atlas_tex.tex, sampler_info),
        };

        let slice = mesh.to_slice(gfx::TriangleList);
        let program = self.graphics.device.link_program(
            VERTEX_SRC.clone(), FRAGMENT_SRC.clone()).unwrap();
        let batch: gfx::batch::RefBatch<_ShaderParamLink, ShaderParam> = self.graphics.make_batch(
            &program, &mesh, slice, &gfx::DrawState::new()).unwrap();
        self.graphics.draw(&batch, &params, &self.frame);
    }

    pub fn draw_line(&mut self, p1: V2<int>, p2: V2<int>, thickness: f32, color: &Rgb) {
        let scale = self.resolution.map(|x| 2.0 / (x as f32));
        let p1 = p1.map(|x| x as f32).mul(scale) - V2(1f32, 1f32);
        let p2 = p2.map(|x| x as f32).mul(scale) - V2(1f32, 1f32);

        let mesh = self.graphics.device.create_mesh([
            Vertex { pos: [p1.0, -p1.1], tex_coord: [0.0, 0.0] },
            Vertex { pos: [p2.0, -p2.1], tex_coord: [0.0, 0.0] },
        ]);

        let sampler_info = Some(self.graphics.device.create_sampler(
            tex::SamplerInfo::new(tex::Scale, tex::Clamp)));
        let params = ShaderParam {
            color: color.to_array(),
            s_texture: (self.flatshade, sampler_info),
        };
        let slice = mesh.to_slice(gfx::Line);
        let program = self.graphics.device.link_program(
            VERTEX_SRC.clone(), FRAGMENT_SRC.clone()).unwrap();
        let mut draw_state = gfx::DrawState::new();
        draw_state.primitive.method = gfx::state::Line(thickness);
        let batch: gfx::batch::RefBatch<_ShaderParamLink, ShaderParam> = self.graphics.make_batch(
            &program, &mesh, slice, &draw_state).unwrap();
        self.graphics.draw(&batch, &params, &self.frame);
    }

    pub fn font_image(&self, c: char) -> Option<Image> {
        let idx = c as int;
        if idx >= 32 && idx < 128 {
            Some(Image((idx - 32) as uint))
        } else {
            None
        }
    }

    pub fn image_dim(&self, Image(idx): Image) -> V2<uint> {
        self.atlas.vertices[idx].1.map(|x| x as uint)
    }
}

pub enum Event<'a> {
    /// Time to render the screen. Call your own render code on the Context
    /// value when you get this.
    Render(&'a mut Context),
    /// Some printable text entered. Data is whole strings to account for
    /// exotic input devices.
    Text(String),
    KeyPressed(key::Key),
    KeyReleased(key::Key),
}

impl<'a> Iterator<Event<'a>> for Context {
    fn next(&mut self) -> Option<Event<'a>> {
        // After a render event, control will return here on a new
        // iter call. Do post-render work here.
        if self.state == EndFrame {
            self.state = Normal;
            self.graphics.end_frame();
            self.window.swap_buffers();
        }

        loop {
            if self.window.should_close() {
                return None;
            }

            self.glfw.poll_events();

            match self.events.try_recv() {
                Ok((_, event)) => {
                    match event {
                        glfw::CharEvent(ch) => {
                            return Some(Text(String::from_char(1, ch)));
                        }
                        glfw::KeyEvent(k, _scan, action, _mods) => {
                            match key::translate_glfw_key(k).map(|k| {
                                if action == glfw::Press || action == glfw::Repeat {
                                    KeyPressed(k)
                                }
                                else {
                                    KeyReleased(k)
                                }
                            }) {
                                Some(e) => { return Some(e); }
                                _ => ()
                            }
                        }
                        // TODO Mouse events.
                        _ => ()
                    }
                }
                _ => ()
            }

            let t = time::precise_time_s();
            if self.frame_interval.map_or(true,
                |x| t - self.last_render_time >= x) {
                self.last_render_time = t;

                // Time to render, must return a handle to self.
                // XXX: Need unsafe hackery to get around lifetimes check.
                self.state = EndFrame;
                let (w, h) = self.window.get_framebuffer_size();
                self.frame = gfx::Frame::new(w as u16, h as u16);
                unsafe {
                    return Some(Render(mem::transmute(self)))
                }
            }
        }
    }
}

/// Drawable images stored in the Canvas.
#[deriving(Clone, PartialEq)]
pub struct Image(uint);

static VERTEX_SRC: gfx::ShaderSource = shaders! {
GLSL_120: b"
    #version 120

    uniform vec4 color;

    attribute vec2 a_pos;
    attribute vec2 a_tex_coord;

    // TODO: Make color a uniform argument.
    varying vec2 v_tex_coord;
    varying vec4 v_color;

    void main() {
        v_tex_coord = a_tex_coord;
        v_color = color;
        gl_Position = vec4(a_pos, 0.0, 1.0);
    }
"
};

static FRAGMENT_SRC: gfx::ShaderSource = shaders! {
GLSL_120: b"
    #version 120

    uniform sampler2D s_texture;

    varying vec2 v_tex_coord;
    varying vec4 v_color;

    void main() {
        vec4 tex_color = texture2D(s_texture, v_tex_coord);
        if (tex_color.a == 0.0) discard;
        gl_FragColor = v_color * tex_color;
    }
"
};

#[shader_param(Program)]
pub struct ShaderParam {
    pub color: [f32, ..4],
    pub s_texture: gfx::shade::TextureParam,
}

#[vertex_format]
struct Vertex {
    #[name = "a_pos"]
    pos: [f32, ..2],

    #[name = "a_tex_coord"]
    tex_coord: [f32, ..2],
}

impl Clone for Vertex {
    fn clone(&self) -> Vertex { *self }
}

struct Texture {
    tex: gfx::TextureHandle,
    pub width: u32,
    pub height: u32,
}

impl Texture {
    fn from_rgba8<D: Device<C>, C: CommandBuffer>(
        img: &ImageBuf<Rgba<u8>>,
        d: &mut D) -> Texture {
        let (w, h) = img.dimensions();
        let mut info = tex::TextureInfo::new();
        info.width = w as u16;
        info.height = h as u16;
        info.kind = tex::Texture2D;
        info.format = tex::RGBA8;

        let tex = d.create_texture(info).unwrap();
        d.update_texture(&tex, &info.to_image_info(), img.pixelbuf()).unwrap();

        Texture {
            tex: tex,
            width: w,
            height: h,
        }
    }
}
<|endoftext|>