<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Utilities for working with sessions.
//!
//! Sessions are created per action request and should be used as the only way
//! to communicate with the server. Sessions are responsible for handling action
//! errors (when something goes wrong) or notifying the server that the action
//! finished (by sending appropriate status signal).
//!
//! They also keep track of various statistics (such as number of transferred
//! bytes, action runtime, etc.) and stop the execution if they exceed limits
//! for a particular request.

mod demand;
mod error;
mod response;
mod sink;

use std::convert::TryInto;

use log::{error, info};

use crate::action;
use crate::message;
pub use self::demand::{Demand, Header, Payload};
pub use self::error::{Error, ParseError, MissingFieldError};
use self::response::{Response, Status};
pub use self::sink::{Sink};

/// A specialized `Result` type for sessions.
pub type Result<T> = std::result::Result<T, Error>;

/// Object associating a session with particular action request.
///
/// This is just a convenience type used to avoid threading large numbers of
/// parameters through different function calls.
///
/// Note that the payload can be quite large object, so it should not be stored
/// with the session itself. The payload should be consumed as soon as the
/// request is dispatched to a particular handler and parsed to a concrete type.
pub struct Task<'s, S: Session> {
    /// A session that this task should be run on.
    pub session: &'s mut S,
    /// Serialized data of the action request.
    pub payload: Payload,
}

impl<'s, S: Session> Task<'s, S> {

    /// Executes the task with a particular action handler.
    pub fn execute<R, H>(self, handler: H) -> Result<()>
    where
        R: action::Request,
        H: FnOnce(&mut S, R) -> Result<()>,
    {
        let request = self.payload.parse()?;
        handler(self.session, request)
    }
}

/// Processes given message, handling all errors.
///
/// This function takes a message from the GRR server and interprets it as an
/// action request. It creates a new session object and dispatches the request
/// to an appropriate action handler. Once the action finishes, the function
/// will send a final status message, informing the server about a success
/// or a failure.
///
/// Note that if action execution fails, this function deals with all the errors
/// by sending appropriate information to the server (if possible), logging them
/// and failing hard if a critical error (e.g. communication failure) occurred.
pub fn handle<M>(message: M)
where
    M: TryInto<Demand, Error=ParseError>,
{
    let demand = match message.try_into() {
        Ok(demand) => {
            info!("requested to execute the '{}' action", demand.action);
            demand
        }
        Err(error) => {
            error!("failed to parse the message: {}", error);
            return;
        }
    };

    let mut session = Action::from_demand(&demand);

    let result = action::dispatch(&demand.action, Task {
        session: &mut session,
        payload: demand.payload,
    });

    if let Err(ref error) = result {
        error!("failed to execute the '{}' action: {}", demand.action, error);
    } else {
        info!("finished executing the '{}' action", demand.action);
    }

    let message = match session.status(result).try_into() {
        Ok(message) => message,
        Err(error) => {
            // If we cannot encode the final status message, there is nothing
            // we can do to notify the server, as status is responsible for
            // reporting errors. We can only log the error and carry on.
            error!("failed to encode status message: {}", error);
            return;
        }
    };

    message::send(message);
}

/// Abstraction for various kinds of sessions.
pub trait Session {
    /// Sends a reply to the flow that call the action.
    fn reply<R>(&mut self, response: R) -> Result<()>
    where R: action::Response + 'static;

    /// Sends a message to a particular sink.
    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where R: action::Response + 'static;

    /// Sends a heartbeat signal to the Fleetspeak process.
    fn heartbeat(&mut self) {
        // TODO: Create a real implementation.
    }
}

/// A session type for unrequested action executions.
///
/// Certain kind of actions are executed not only when a server flow decides to
/// do so, but also upon particular kind of events (e.g. the agent's startup).
/// In such cases, when one needs to trigger action execution manually, ad-hoc
/// sessions should be used.
pub struct Adhoc;

impl Session for Adhoc {

    // TODO: Session trait should be probably split into two traits and then
    // make the actions that do not care about the `reply` method implement the
    // simpler one.
    fn reply<R>(&mut self, response: R) -> Result<()>
    where
        R: action::Response,
    {
        error!("attempted to reply to an ad-hoc session, dropping response");
        drop(response);

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where
        R: action::Response,
    {
        send(sink.wrap(response))?;

        Ok(())
    }
}

/// A session type for ordinary action requests.
///
/// This is a normal session type that that is associated with some flow on the
/// server. It keeps track of the responses it sends and collects statistics
/// about network and runtime utilization to kill the action if it is needed.
pub struct Action {
    header: Header,
    next_response_id: u64,
}

impl Action {

    /// Constructs a new session for the given `demand` object.
    pub fn from_demand(demand: &Demand) -> Action {
        // Response identifiers that GRR agents use start at 1. Unfortunately,
        // the server uses this assumption (to determine the number of expected
        // responses when status message is received), so we have to follow this
        // behaviour in RRG as well.
        Action {
            header: demand.header.clone(),
            next_response_id: 1,
        }
    }

    /// Wraps an action response to a session-specific response.
    fn wrap<R>(&self, response: R) -> Response<R>
    where
        R: action::Response
    {
        Response {
            session_id: self.header.session_id.clone(),
            request_id: Some(self.header.request_id),
            response_id: Some(self.next_response_id),
            data: response,
        }
    }

    /// Wraps an action result to a session-specific status response.
    ///
    /// Note that this method consumes the session. The reason for this is that
    /// status response should be obtained as the very last response and no
    /// further replies should be possible after that.
    fn status(self, result: Result<()>) -> Status {
        Status {
            session_id: self.header.session_id,
            request_id: self.header.request_id,
            response_id: self.next_response_id,
            result: result,
        }
    }
}

impl Session for Action {

    fn reply<R: action::Response>(&mut self, response: R) -> Result<()> {
        send(self.wrap(response))?;
        self.next_response_id += 1;

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where
        R: action::Response,
    {
        send(sink.wrap(response))?;

        Ok(())
    }
}

/// Sends a session response to the server.
///
/// Note that this function is not exposed on purpose. Actions should send
/// responses through session objects which introduce a layer of safety. `send`
/// is a low-level utility supposed to be used internally.
fn send<R>(response: Response<R>) -> Result<()>
where
    R: action::Response,
{
    let message = response.try_into()?;
    message::send(message);

    Ok(())
}

#[cfg(test)]
pub mod test {
    use std::any::Any;
    use std::collections::HashMap;

    use super::*;

    /// A session type intended to be used in tests.
    ///
    /// Testing actions with normal session objects can be quite hard, since
    /// they communicate with the outside world (through Fleetspeak). Since we
    /// want to keep the tests minimal and not waste resources on unneeded I/O,
    /// using real sessions is not an option.
    ///
    /// Instead, one can use a `Fake` session. It simply accumulates responses
    /// that the action sends and lets the creator inspect them later.
    pub struct Fake {
        replies: Vec<Box<dyn Any>>,
        responses: HashMap<Sink, Vec<Box<dyn Any>>>,
    }

    impl Fake {

        /// Constructs a new fake session.
        pub fn new() -> Fake {
            Fake {
                replies: Vec::new(),
                responses: std::collections::HashMap::new(),
            }
        }

        /// Yields the number of replies that this session sent so far.
        pub fn reply_count(&self) -> usize {
            self.replies.len()
        }

        /// Retrieves a reply corresponding to the given id.
        ///
        /// The identifier corresponding to the first response is 0, the second
        /// one is 1 and so on.
        ///
        /// This method will panic if a reply with the specified `id` does not
        /// exist or if it exists but has a wrong type.
        pub fn reply<R>(&self, id: usize) -> &R
        where
            R: action::Response + 'static,
        {
            let reply = match self.replies.get(id) {
                Some(reply) => reply,
                None => panic!("no reply #{}", id),
            };

            reply.downcast_ref().expect("unexpected reply type")
        }

        /// Yields the number of responses sent so far to the specified sink.
        pub fn response_count(&self, sink: Sink) -> usize {
            match self.responses.get(&sink) {
                Some(responses) => responses.len(),
                None => 0,
            }
        }

        /// Retrieves a response with the given id sent to a particular sink.
        ///
        /// The identifier corresponding to the first response to the particular
        /// sink is 0, to the second one (to the same sink) is 1 and so on.
        ///
        /// This method will panic if a reply with the specified `id` to the
        /// given `sink` does not exist or if it exists but has wrong type.
        pub fn response<R>(&self, sink: Sink, id: usize) -> &R
        where
            R: action::Response + 'static,
        {
            let responses = match self.responses.get(&sink) {
                Some(responses) => responses,
                None => panic!("no responses for sink '{:?}'", sink),
            };

            let response = match responses.get(id) {
                Some(response) => response,
                None => panic!("no response #{} for sink '{:?}'", id, sink),
            };

            match response.downcast_ref() {
                Some(response) => response,
                None => panic!("unexpected response type in sink '{:?}'", sink),
            }
        }
    }

    impl Session for Fake {

        fn reply<R>(&mut self, response: R) -> Result<()>
        where
            R: action::Response + 'static,
        {
            self.replies.push(Box::new(response));

            Ok(())
        }

        fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
        where
            R: action::Response + 'static,
        {
            let responses = self.responses.entry(sink).or_insert_with(Vec::new);
            responses.push(Box::new(response));

            Ok(())
        }
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn test_fake_reply_count() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.reply(()).unwrap();
            session.reply(()).unwrap();
            session.reply(()).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        assert_eq!(session.reply_count(), 3);
    }

    #[test]
    fn test_fake_response_count() {

        // TODO: Extend this test with more sinks (once we have some more sinks
        // defined).

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, ()).unwrap();
            session.send(Sink::STARTUP, ()).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        assert_eq!(session.response_count(Sink::STARTUP), 2);
    }

    #[test]
    fn test_fake_reply_correct_response() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.reply(StringResponse::from("foo")).unwrap();
            session.reply(StringResponse::from("bar")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        assert_eq!(session.reply::<StringResponse>(0).0, "foo");
        assert_eq!(session.reply::<StringResponse>(1).0, "bar");
    }

    #[test]
    #[should_panic(expected = "no reply #0")]
    fn test_fake_reply_incorrect_response_id() {

        fn handle<S: Session>(_: &mut S, _: ()) {
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.reply::<()>(0);
    }

    #[test]
    #[should_panic(expected = "unexpected reply type")]
    fn test_fake_reply_incorrect_response_type() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.reply(StringResponse::from("quux")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.reply::<()>(0);
    }

    #[test]
    fn test_fake_response_correct_response() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, StringResponse::from("foo")).unwrap();
            session.send(Sink::STARTUP, StringResponse::from("bar")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        let response_foo = session.response::<StringResponse>(Sink::STARTUP, 0);
        let response_bar = session.response::<StringResponse>(Sink::STARTUP, 1);
        assert_eq!(response_foo.0, "foo");
        assert_eq!(response_bar.0, "bar");
    }

    #[test]
    #[should_panic(expected = "no responses")]
    fn test_fake_response_empty_sink() {

        fn handle<S: Session>(_: &mut S, _: ()) {
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.response::<()>(Sink::STARTUP, 0);
    }

    #[test]
    #[should_panic(expected = "no response #42")]
    fn test_fake_response_incorrect_response_id() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, ()).unwrap();
            session.send(Sink::STARTUP, ()).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.response::<()>(Sink::STARTUP, 42);
    }

    #[test]
    #[should_panic(expected = "unexpected response type")]
    fn test_fake_response_incorrect_response_type() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, StringResponse::from("quux")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.response::<()>(Sink::STARTUP, 0);
    }

    struct StringResponse(String);

    impl<S: Into<String>> From<S> for StringResponse {

        fn from(string: S) -> StringResponse {
            StringResponse(string.into())
        }
    }

    impl action::Response for StringResponse {

        const RDF_NAME: Option<&'static str> = Some("RDFString");

        type Proto = String;

        fn into_proto(self) -> String {
            self.0
        }
    }
}
<commit_msg>Implement an iterator over fake session replies.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Utilities for working with sessions.
//!
//! Sessions are created per action request and should be used as the only way
//! to communicate with the server. Sessions are responsible for handling action
//! errors (when something goes wrong) or notifying the server that the action
//! finished (by sending appropriate status signal).
//!
//! They also keep track of various statistics (such as number of transferred
//! bytes, action runtime, etc.) and stop the execution if they exceed limits
//! for a particular request.

mod demand;
mod error;
mod response;
mod sink;

use std::convert::TryInto;

use log::{error, info};

use crate::action;
use crate::message;
pub use self::demand::{Demand, Header, Payload};
pub use self::error::{Error, ParseError, MissingFieldError};
use self::response::{Response, Status};
pub use self::sink::{Sink};

/// A specialized `Result` type for sessions.
pub type Result<T> = std::result::Result<T, Error>;

/// Object associating a session with particular action request.
///
/// This is just a convenience type used to avoid threading large numbers of
/// parameters through different function calls.
///
/// Note that the payload can be quite large object, so it should not be stored
/// with the session itself. The payload should be consumed as soon as the
/// request is dispatched to a particular handler and parsed to a concrete type.
pub struct Task<'s, S: Session> {
    /// A session that this task should be run on.
    pub session: &'s mut S,
    /// Serialized data of the action request.
    pub payload: Payload,
}

impl<'s, S: Session> Task<'s, S> {

    /// Executes the task with a particular action handler.
    pub fn execute<R, H>(self, handler: H) -> Result<()>
    where
        R: action::Request,
        H: FnOnce(&mut S, R) -> Result<()>,
    {
        let request = self.payload.parse()?;
        handler(self.session, request)
    }
}

/// Processes given message, handling all errors.
///
/// This function takes a message from the GRR server and interprets it as an
/// action request. It creates a new session object and dispatches the request
/// to an appropriate action handler. Once the action finishes, the function
/// will send a final status message, informing the server about a success
/// or a failure.
///
/// Note that if action execution fails, this function deals with all the errors
/// by sending appropriate information to the server (if possible), logging them
/// and failing hard if a critical error (e.g. communication failure) occurred.
pub fn handle<M>(message: M)
where
    M: TryInto<Demand, Error=ParseError>,
{
    let demand = match message.try_into() {
        Ok(demand) => {
            info!("requested to execute the '{}' action", demand.action);
            demand
        }
        Err(error) => {
            error!("failed to parse the message: {}", error);
            return;
        }
    };

    let mut session = Action::from_demand(&demand);

    let result = action::dispatch(&demand.action, Task {
        session: &mut session,
        payload: demand.payload,
    });

    if let Err(ref error) = result {
        error!("failed to execute the '{}' action: {}", demand.action, error);
    } else {
        info!("finished executing the '{}' action", demand.action);
    }

    let message = match session.status(result).try_into() {
        Ok(message) => message,
        Err(error) => {
            // If we cannot encode the final status message, there is nothing
            // we can do to notify the server, as status is responsible for
            // reporting errors. We can only log the error and carry on.
            error!("failed to encode status message: {}", error);
            return;
        }
    };

    message::send(message);
}

/// Abstraction for various kinds of sessions.
pub trait Session {
    /// Sends a reply to the flow that call the action.
    fn reply<R>(&mut self, response: R) -> Result<()>
    where R: action::Response + 'static;

    /// Sends a message to a particular sink.
    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where R: action::Response + 'static;

    /// Sends a heartbeat signal to the Fleetspeak process.
    fn heartbeat(&mut self) {
        // TODO: Create a real implementation.
    }
}

/// A session type for unrequested action executions.
///
/// Certain kind of actions are executed not only when a server flow decides to
/// do so, but also upon particular kind of events (e.g. the agent's startup).
/// In such cases, when one needs to trigger action execution manually, ad-hoc
/// sessions should be used.
pub struct Adhoc;

impl Session for Adhoc {

    // TODO: Session trait should be probably split into two traits and then
    // make the actions that do not care about the `reply` method implement the
    // simpler one.
    fn reply<R>(&mut self, response: R) -> Result<()>
    where
        R: action::Response,
    {
        error!("attempted to reply to an ad-hoc session, dropping response");
        drop(response);

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where
        R: action::Response,
    {
        send(sink.wrap(response))?;

        Ok(())
    }
}

/// A session type for ordinary action requests.
///
/// This is a normal session type that that is associated with some flow on the
/// server. It keeps track of the responses it sends and collects statistics
/// about network and runtime utilization to kill the action if it is needed.
pub struct Action {
    header: Header,
    next_response_id: u64,
}

impl Action {

    /// Constructs a new session for the given `demand` object.
    pub fn from_demand(demand: &Demand) -> Action {
        // Response identifiers that GRR agents use start at 1. Unfortunately,
        // the server uses this assumption (to determine the number of expected
        // responses when status message is received), so we have to follow this
        // behaviour in RRG as well.
        Action {
            header: demand.header.clone(),
            next_response_id: 1,
        }
    }

    /// Wraps an action response to a session-specific response.
    fn wrap<R>(&self, response: R) -> Response<R>
    where
        R: action::Response
    {
        Response {
            session_id: self.header.session_id.clone(),
            request_id: Some(self.header.request_id),
            response_id: Some(self.next_response_id),
            data: response,
        }
    }

    /// Wraps an action result to a session-specific status response.
    ///
    /// Note that this method consumes the session. The reason for this is that
    /// status response should be obtained as the very last response and no
    /// further replies should be possible after that.
    fn status(self, result: Result<()>) -> Status {
        Status {
            session_id: self.header.session_id,
            request_id: self.header.request_id,
            response_id: self.next_response_id,
            result: result,
        }
    }
}

impl Session for Action {

    fn reply<R: action::Response>(&mut self, response: R) -> Result<()> {
        send(self.wrap(response))?;
        self.next_response_id += 1;

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where
        R: action::Response,
    {
        send(sink.wrap(response))?;

        Ok(())
    }
}

/// Sends a session response to the server.
///
/// Note that this function is not exposed on purpose. Actions should send
/// responses through session objects which introduce a layer of safety. `send`
/// is a low-level utility supposed to be used internally.
fn send<R>(response: Response<R>) -> Result<()>
where
    R: action::Response,
{
    let message = response.try_into()?;
    message::send(message);

    Ok(())
}

#[cfg(test)]
pub mod test {
    use std::any::Any;
    use std::collections::HashMap;

    use super::*;

    /// A session type intended to be used in tests.
    ///
    /// Testing actions with normal session objects can be quite hard, since
    /// they communicate with the outside world (through Fleetspeak). Since we
    /// want to keep the tests minimal and not waste resources on unneeded I/O,
    /// using real sessions is not an option.
    ///
    /// Instead, one can use a `Fake` session. It simply accumulates responses
    /// that the action sends and lets the creator inspect them later.
    pub struct Fake {
        replies: Vec<Box<dyn Any>>,
        responses: HashMap<Sink, Vec<Box<dyn Any>>>,
    }

    impl Fake {

        /// Constructs a new fake session.
        pub fn new() -> Fake {
            Fake {
                replies: Vec::new(),
                responses: std::collections::HashMap::new(),
            }
        }

        /// Yields the number of replies that this session sent so far.
        pub fn reply_count(&self) -> usize {
            self.replies.len()
        }

        /// Retrieves a reply corresponding to the given id.
        ///
        /// The identifier corresponding to the first response is 0, the second
        /// one is 1 and so on.
        ///
        /// This method will panic if a reply with the specified `id` does not
        /// exist or if it exists but has a wrong type.
        pub fn reply<R>(&self, id: usize) -> &R
        where
            R: action::Response + 'static,
        {
            match self.replies::<R>().nth(id) {
                Some(reply) => reply,
                None => panic!("no reply #{}", id),
            }
        }

        pub fn replies<R>(&self) -> impl Iterator<Item = &R>
        where
            R: action::Response + 'static
        {
            self.replies.iter().map(|reply| {
                reply.downcast_ref().expect("unexpected reply type")
            })
        }

        /// Yields the number of responses sent so far to the specified sink.
        pub fn response_count(&self, sink: Sink) -> usize {
            match self.responses.get(&sink) {
                Some(responses) => responses.len(),
                None => 0,
            }
        }

        /// Retrieves a response with the given id sent to a particular sink.
        ///
        /// The identifier corresponding to the first response to the particular
        /// sink is 0, to the second one (to the same sink) is 1 and so on.
        ///
        /// This method will panic if a reply with the specified `id` to the
        /// given `sink` does not exist or if it exists but has wrong type.
        pub fn response<R>(&self, sink: Sink, id: usize) -> &R
        where
            R: action::Response + 'static,
        {
            let responses = match self.responses.get(&sink) {
                Some(responses) => responses,
                None => panic!("no responses for sink '{:?}'", sink),
            };

            let response = match responses.get(id) {
                Some(response) => response,
                None => panic!("no response #{} for sink '{:?}'", id, sink),
            };

            match response.downcast_ref() {
                Some(response) => response,
                None => panic!("unexpected response type in sink '{:?}'", sink),
            }
        }
    }

    impl Session for Fake {

        fn reply<R>(&mut self, response: R) -> Result<()>
        where
            R: action::Response + 'static,
        {
            self.replies.push(Box::new(response));

            Ok(())
        }

        fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
        where
            R: action::Response + 'static,
        {
            let responses = self.responses.entry(sink).or_insert_with(Vec::new);
            responses.push(Box::new(response));

            Ok(())
        }
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn test_fake_reply_count() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.reply(()).unwrap();
            session.reply(()).unwrap();
            session.reply(()).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        assert_eq!(session.reply_count(), 3);
    }

    #[test]
    fn test_fake_response_count() {

        // TODO: Extend this test with more sinks (once we have some more sinks
        // defined).

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, ()).unwrap();
            session.send(Sink::STARTUP, ()).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        assert_eq!(session.response_count(Sink::STARTUP), 2);
    }

    #[test]
    fn test_fake_reply_correct_response() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.reply(StringResponse::from("foo")).unwrap();
            session.reply(StringResponse::from("bar")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        assert_eq!(session.reply::<StringResponse>(0).0, "foo");
        assert_eq!(session.reply::<StringResponse>(1).0, "bar");
    }

    #[test]
    #[should_panic(expected = "no reply #0")]
    fn test_fake_reply_incorrect_response_id() {

        fn handle<S: Session>(_: &mut S, _: ()) {
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.reply::<()>(0);
    }

    #[test]
    #[should_panic(expected = "unexpected reply type")]
    fn test_fake_reply_incorrect_response_type() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.reply(StringResponse::from("quux")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.reply::<()>(0);
    }

    #[test]
    fn test_fake_response_correct_response() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, StringResponse::from("foo")).unwrap();
            session.send(Sink::STARTUP, StringResponse::from("bar")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        let response_foo = session.response::<StringResponse>(Sink::STARTUP, 0);
        let response_bar = session.response::<StringResponse>(Sink::STARTUP, 1);
        assert_eq!(response_foo.0, "foo");
        assert_eq!(response_bar.0, "bar");
    }

    #[test]
    #[should_panic(expected = "no responses")]
    fn test_fake_response_empty_sink() {

        fn handle<S: Session>(_: &mut S, _: ()) {
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.response::<()>(Sink::STARTUP, 0);
    }

    #[test]
    #[should_panic(expected = "no response #42")]
    fn test_fake_response_incorrect_response_id() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, ()).unwrap();
            session.send(Sink::STARTUP, ()).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.response::<()>(Sink::STARTUP, 42);
    }

    #[test]
    #[should_panic(expected = "unexpected response type")]
    fn test_fake_response_incorrect_response_type() {

        fn handle<S: Session>(session: &mut S, _: ()) {
            session.send(Sink::STARTUP, StringResponse::from("quux")).unwrap();
        }

        let mut session = test::Fake::new();
        handle(&mut session, ());

        session.response::<()>(Sink::STARTUP, 0);
    }

    struct StringResponse(String);

    impl<S: Into<String>> From<S> for StringResponse {

        fn from(string: S) -> StringResponse {
            StringResponse(string.into())
        }
    }

    impl action::Response for StringResponse {

        const RDF_NAME: Option<&'static str> = Some("RDFString");

        type Proto = String;

        fn into_proto(self) -> String {
            self.0
        }
    }
}
<|endoftext|>


<commit_before>use camera::Camera;
use event::{Action, Key, MouseButton, WindowEvent};
use na::{self, Isometry3, Matrix4, Perspective3, Point3, Vector2, Vector3};
use resource::ShaderUniform;
use std::f32;
use window::Canvas;

/// Arc-ball camera mode.
///
/// An arc-ball camera is a camera rotating around a fixed point (the focus point) and always
/// looking at it. The following inputs are handled:
///
/// * Left button press + drag - rotates the camera around the focus point
/// * Right button press + drag - translates the focus point on the plane orthogonal to the view
/// direction
/// * Scroll in/out - zoom in/out
/// * Enter key - set the focus point to the origin
#[derive(Clone, Debug)]
pub struct ArcBall {
    /// The focus point.
    at: Point3<f32>,
    /// Yaw of the camera (rotation along the y axis).
    yaw: f32,
    /// Pitch of the camera (rotation along the x axis).
    pitch: f32,
    /// Distance from the camera to the `at` focus point.
    dist: f32,
    /// Minimum distance from the camera to the `at` focus point.
    min_dist: f32,
    /// Maximum distance from the camera to the `at` focus point.
    max_dist: f32,

    /// Increment of the yaw per unit mouse movement. The default value is 0.005.
    yaw_step: f32,
    /// Increment of the pitch per unit mouse movement. The default value is 0.005.
    pitch_step: f32,
    /// Minimum pitch of the camera.
    min_pitch: f32,
    /// Maximum pitch ofthe camera.
    max_pitch: f32,
    /// Increment of the distance per unit scrolling. The default value is 40.0.
    dist_step: f32,
    rotate_button: Option<MouseButton>,
    drag_button: Option<MouseButton>,
    reset_key: Option<Key>,

    projection: Perspective3<f32>,
    view: Matrix4<f32>,
    proj: Matrix4<f32>,
    proj_view: Matrix4<f32>,
    inverse_proj_view: Matrix4<f32>,
    last_cursor_pos: Vector2<f32>,
    up_axis: Vector3<f32>,
}

impl ArcBall {
    /// Create a new arc-ball camera.
    pub fn new(eye: Point3<f32>, at: Point3<f32>) -> ArcBall {
        ArcBall::new_with_frustrum(f32::consts::PI / 4.0, 0.1, 1024.0, eye, at)
    }

    /// Creates a new arc ball camera with default sensitivity values.
    pub fn new_with_frustrum(
        fov: f32,
        znear: f32,
        zfar: f32,
        eye: Point3<f32>,
        at: Point3<f32>,
    ) -> ArcBall {
        let mut res = ArcBall {
            at: Point3::new(0.0, 0.0, 0.0),
            yaw: 0.0,
            pitch: 0.0,
            dist: 0.0,
            min_dist: 0.00001,
            max_dist: std::f32::MAX,
            yaw_step: 0.005,
            pitch_step: 0.005,
            min_pitch: 0.01,
            max_pitch: std::f32::consts::PI - 0.01,
            dist_step: 40.0,
            rotate_button: Some(MouseButton::Button1),
            drag_button: Some(MouseButton::Button2),
            reset_key: Some(Key::Return),
            projection: Perspective3::new(800.0 / 600.0, fov, znear, zfar),
            view: na::zero(),
            proj: na::zero(),
            proj_view: na::zero(),
            inverse_proj_view: na::zero(),
            last_cursor_pos: na::zero(),
            up_axis: Vector3::y(),
        };

        res.look_at(eye, at);

        res
    }

    /// The point the arc-ball is looking at.
    pub fn at(&self) -> Point3<f32> {
        self.at
    }

    /// Get a mutable reference to the point the camera is looking at.
    pub fn set_at(&mut self, at: Point3<f32>) {
        self.at = at;
        self.update_projviews();
    }

    /// The arc-ball camera `yaw`.
    pub fn yaw(&self) -> f32 {
        self.yaw
    }

    /// Sets the camera `yaw`. Change this to modify the rotation along the `up` axis.
    pub fn set_yaw(&mut self, yaw: f32) {
        self.yaw = yaw;

        self.update_restrictions();
        self.update_projviews();
    }

    /// The arc-ball camera `pitch`.
    pub fn pitch(&self) -> f32 {
        self.pitch
    }

    /// Sets the camera `pitch`.
    pub fn set_pitch(&mut self, pitch: f32) {
        self.pitch = pitch;

        self.update_restrictions();
        self.update_projviews();
    }

    /// The minimum pitch of the camera.
    pub fn min_pitch(&self) -> f32 {
        self.min_pitch
    }

    /// Set the minimum pitch of the camera.
    pub fn set_min_pitch(&mut self, min_pitch: f32) {
        self.min_pitch = min_pitch;
    }

    /// The maximum pitch of the camera.
    pub fn max_pitch(&self) -> f32 {
        self.max_pitch
    }

    /// Set the maximum pitch of the camera.
    pub fn set_max_pitch(&mut self, max_pitch: f32) {
        self.max_pitch = max_pitch;
    }

    /// The distance from the camera position to its view point.
    pub fn dist(&self) -> f32 {
        self.dist
    }

    /// Move the camera such that it is at a given distance from the view point.
    pub fn set_dist(&mut self, dist: f32) {
        self.dist = dist;

        self.update_restrictions();
        self.update_projviews();
    }

    /// The minimum distance from the camera position to its view point.
    pub fn min_dist(&self) -> f32 {
        self.min_dist
    }

    /// Set the minimum distance from the camera position to its view point.
    pub fn set_min_dist(&mut self, min_dist: f32) {
        self.min_dist = min_dist;
    }

    /// The maximum distance from the camera position to its view point.
    pub fn max_dist(&self) -> f32 {
        self.max_dist
    }

    /// Set the maximum distance from the camera position to its view point.
    pub fn set_max_dist(&mut self, max_dist: f32) {
        self.max_dist = max_dist;
    }

    /// Move and orient the camera such that it looks at a specific point.
    pub fn look_at(&mut self, eye: Point3<f32>, at: Point3<f32>) {
        let dist = (eye - at).norm();
        let pitch = ((eye.y - at.y) / dist).acos();
        let yaw = (eye.z - at.z).atan2(eye.x - at.x);

        self.at = at;
        self.dist = dist;
        self.yaw = yaw;
        self.pitch = pitch;

        self.update_restrictions();
        self.update_projviews();
    }

    /// Transformation applied by the camera without perspective.
    fn update_restrictions(&mut self) {
        if self.dist < self.min_dist {
            self.dist = self.min_dist
        }

        if self.dist > self.max_dist {
            self.dist = self.max_dist
        }

        if self.pitch <= self.min_pitch {
            self.pitch = self.min_pitch
        }

        if self.pitch > self.max_pitch {
            self.pitch = self.max_pitch
        }
    }

    /// The button used to rotate the ArcBall camera.
    pub fn rotate_button(&self) -> Option<MouseButton> {
        self.rotate_button
    }

    /// Set the button used to rotate the ArcBall camera.
    /// Use None to disable rotation.
    pub fn rebind_rotate_button(&mut self, new_button: Option<MouseButton>) {
        self.rotate_button = new_button;
    }

    /// The button used to drag the ArcBall camera.
    pub fn drag_button(&self) -> Option<MouseButton> {
        self.drag_button
    }

    /// Set the button used to drag the ArcBall camera.
    /// Use None to disable dragging.
    pub fn rebind_drag_button(&mut self, new_button: Option<MouseButton>) {
        self.drag_button = new_button;
    }

    /// The key used to reset the ArcBall camera.
    pub fn reset_key(&self) -> Option<Key> {
        self.reset_key
    }

    /// Set the key used to reset the ArcBall camera.
    /// Use None to disable reset.
    pub fn rebind_reset_key(&mut self, new_key: Option<Key>) {
        self.reset_key = new_key;
    }

    fn handle_left_button_displacement(&mut self, dpos: &Vector2<f32>) {
        self.yaw = self.yaw + dpos.x * self.yaw_step;
        self.pitch = self.pitch - dpos.y * self.pitch_step;

        self.update_restrictions();
        self.update_projviews();
    }

    fn handle_right_button_displacement(&mut self, dpos: &Vector2<f32>) {
        let eye = self.eye();
        let dir = (self.at - eye).normalize();
        let tangent = Vector3::y().cross(&dir).normalize();
        let bitangent = dir.cross(&tangent);
        let mult = self.dist / 1000.0;

        self.at = self.at + tangent * (dpos.x * mult) + bitangent * (dpos.y * mult);
        self.update_projviews();
    }

    fn handle_scroll(&mut self, off: f32) {
        self.dist = self.dist + self.dist_step * (off) / 120.0;
        self.update_restrictions();
        self.update_projviews();
    }

    fn update_projviews(&mut self) {
        self.proj = *self.projection.as_matrix();
        self.view = self.view_transform().to_homogeneous();
        self.proj_view = self.proj * self.view;
        self.inverse_proj_view = self.proj_view.try_inverse().unwrap();
    }

    /// Sets the up vector of this camera.
    #[inline]
    pub fn set_up_axis(&mut self, up_axis: Vector3<f32>) {
        self.up_axis = up_axis;
    }
}

impl Camera for ArcBall {
    fn clip_planes(&self) -> (f32, f32) {
        (self.projection.znear(), self.projection.zfar())
    }

    fn view_transform(&self) -> Isometry3<f32> {
        Isometry3::look_at_rh(&self.eye(), &self.at, &self.up_axis)
    }

    fn eye(&self) -> Point3<f32> {
        let px = self.at.x + self.dist * self.yaw.cos() * self.pitch.sin();
        let py = self.at.y + self.dist * self.pitch.cos();
        let pz = self.at.z + self.dist * self.yaw.sin() * self.pitch.sin();

        Point3::new(px, py, pz)
    }

    fn handle_event(&mut self, canvas: &Canvas, event: &WindowEvent) {
        match *event {
            WindowEvent::CursorPos(x, y, _) => {
                let curr_pos = Vector2::new(x as f32, y as f32);

                if let Some(rotate_button) = self.rotate_button {
                    if canvas.get_mouse_button(rotate_button) == Action::Press {
                        let dpos = curr_pos - self.last_cursor_pos;
                        self.handle_left_button_displacement(&dpos)
                    }
                }

                if let Some(drag_button) = self.drag_button {
                    if canvas.get_mouse_button(drag_button) == Action::Press {
                        let dpos = curr_pos - self.last_cursor_pos;
                        self.handle_right_button_displacement(&dpos)
                    }
                }

                self.last_cursor_pos = curr_pos;
            }
            WindowEvent::Key(key, Action::Press, _) if Some(key) == self.reset_key => {
                self.at = Point3::origin();
                self.update_projviews();
            }
            WindowEvent::Scroll(_, off, _) => self.handle_scroll(off as f32),
            WindowEvent::FramebufferSize(w, h) => {
                self.projection.set_aspect(w as f32 / h as f32);
                self.update_projviews();
            }
            _ => {}
        }
    }

    #[inline]
    fn upload(
        &self,
        _: usize,
        proj: &mut ShaderUniform<Matrix4<f32>>,
        view: &mut ShaderUniform<Matrix4<f32>>,
    ) {
        proj.upload(&self.proj);
        view.upload(&self.view);
    }

    fn transformation(&self) -> Matrix4<f32> {
        self.proj_view
    }

    fn inverse_transformation(&self) -> Matrix4<f32> {
        self.inverse_proj_view
    }

    fn update(&mut self, _: &Canvas) {}
}
<commit_msg>ArcBall camera: add the possibility to filter movement inputs using modifiers.<commit_after>use camera::Camera;
use event::{Action, Key, MouseButton, WindowEvent, Modifiers};
use na::{self, Isometry3, Matrix4, Perspective3, Point3, Vector2, Vector3};
use resource::ShaderUniform;
use std::f32;
use window::Canvas;

/// Arc-ball camera mode.
///
/// An arc-ball camera is a camera rotating around a fixed point (the focus point) and always
/// looking at it. The following inputs are handled:
///
/// * Left button press + drag - rotates the camera around the focus point
/// * Right button press + drag - translates the focus point on the plane orthogonal to the view
/// direction
/// * Scroll in/out - zoom in/out
/// * Enter key - set the focus point to the origin
#[derive(Clone, Debug)]
pub struct ArcBall {
    /// The focus point.
    at: Point3<f32>,
    /// Yaw of the camera (rotation along the y axis).
    yaw: f32,
    /// Pitch of the camera (rotation along the x axis).
    pitch: f32,
    /// Distance from the camera to the `at` focus point.
    dist: f32,
    /// Minimum distance from the camera to the `at` focus point.
    min_dist: f32,
    /// Maximum distance from the camera to the `at` focus point.
    max_dist: f32,

    /// Increment of the yaw per unit mouse movement. The default value is 0.005.
    yaw_step: f32,
    /// Increment of the pitch per unit mouse movement. The default value is 0.005.
    pitch_step: f32,
    /// Minimum pitch of the camera.
    min_pitch: f32,
    /// Maximum pitch ofthe camera.
    max_pitch: f32,
    /// Increment of the distance per unit scrolling. The default value is 40.0.
    dist_step: f32,
    rotate_button: Option<MouseButton>,
    rotate_modifiers: Option<Modifiers>,
    drag_button: Option<MouseButton>,
    drag_modifiers: Option<Modifiers>,
    reset_key: Option<Key>,

    projection: Perspective3<f32>,
    view: Matrix4<f32>,
    proj: Matrix4<f32>,
    proj_view: Matrix4<f32>,
    inverse_proj_view: Matrix4<f32>,
    last_cursor_pos: Vector2<f32>,
    up_axis: Vector3<f32>,
}

impl ArcBall {
    /// Create a new arc-ball camera.
    pub fn new(eye: Point3<f32>, at: Point3<f32>) -> ArcBall {
        ArcBall::new_with_frustrum(f32::consts::PI / 4.0, 0.1, 1024.0, eye, at)
    }

    /// Creates a new arc ball camera with default sensitivity values.
    pub fn new_with_frustrum(
        fov: f32,
        znear: f32,
        zfar: f32,
        eye: Point3<f32>,
        at: Point3<f32>,
    ) -> ArcBall {
        let mut res = ArcBall {
            at: Point3::new(0.0, 0.0, 0.0),
            yaw: 0.0,
            pitch: 0.0,
            dist: 0.0,
            min_dist: 0.00001,
            max_dist: std::f32::MAX,
            yaw_step: 0.005,
            pitch_step: 0.005,
            min_pitch: 0.01,
            max_pitch: std::f32::consts::PI - 0.01,
            dist_step: 40.0,
            rotate_button: Some(MouseButton::Button1),
            rotate_modifiers: None,
            drag_button: Some(MouseButton::Button2),
            drag_modifiers: None,
            reset_key: Some(Key::Return),
            projection: Perspective3::new(800.0 / 600.0, fov, znear, zfar),
            view: na::zero(),
            proj: na::zero(),
            proj_view: na::zero(),
            inverse_proj_view: na::zero(),
            last_cursor_pos: na::zero(),
            up_axis: Vector3::y(),
        };

        res.look_at(eye, at);

        res
    }

    /// The point the arc-ball is looking at.
    pub fn at(&self) -> Point3<f32> {
        self.at
    }

    /// Get a mutable reference to the point the camera is looking at.
    pub fn set_at(&mut self, at: Point3<f32>) {
        self.at = at;
        self.update_projviews();
    }

    /// The arc-ball camera `yaw`.
    pub fn yaw(&self) -> f32 {
        self.yaw
    }

    /// Sets the camera `yaw`. Change this to modify the rotation along the `up` axis.
    pub fn set_yaw(&mut self, yaw: f32) {
        self.yaw = yaw;

        self.update_restrictions();
        self.update_projviews();
    }

    /// The arc-ball camera `pitch`.
    pub fn pitch(&self) -> f32 {
        self.pitch
    }

    /// Sets the camera `pitch`.
    pub fn set_pitch(&mut self, pitch: f32) {
        self.pitch = pitch;

        self.update_restrictions();
        self.update_projviews();
    }

    /// The minimum pitch of the camera.
    pub fn min_pitch(&self) -> f32 {
        self.min_pitch
    }

    /// Set the minimum pitch of the camera.
    pub fn set_min_pitch(&mut self, min_pitch: f32) {
        self.min_pitch = min_pitch;
    }

    /// The maximum pitch of the camera.
    pub fn max_pitch(&self) -> f32 {
        self.max_pitch
    }

    /// Set the maximum pitch of the camera.
    pub fn set_max_pitch(&mut self, max_pitch: f32) {
        self.max_pitch = max_pitch;
    }

    /// The distance from the camera position to its view point.
    pub fn dist(&self) -> f32 {
        self.dist
    }

    /// Move the camera such that it is at a given distance from the view point.
    pub fn set_dist(&mut self, dist: f32) {
        self.dist = dist;

        self.update_restrictions();
        self.update_projviews();
    }

    /// The minimum distance from the camera position to its view point.
    pub fn min_dist(&self) -> f32 {
        self.min_dist
    }

    /// Set the minimum distance from the camera position to its view point.
    pub fn set_min_dist(&mut self, min_dist: f32) {
        self.min_dist = min_dist;
    }

    /// The maximum distance from the camera position to its view point.
    pub fn max_dist(&self) -> f32 {
        self.max_dist
    }

    /// Set the maximum distance from the camera position to its view point.
    pub fn set_max_dist(&mut self, max_dist: f32) {
        self.max_dist = max_dist;
    }

    /// Move and orient the camera such that it looks at a specific point.
    pub fn look_at(&mut self, eye: Point3<f32>, at: Point3<f32>) {
        let dist = (eye - at).norm();
        let pitch = ((eye.y - at.y) / dist).acos();
        let yaw = (eye.z - at.z).atan2(eye.x - at.x);

        self.at = at;
        self.dist = dist;
        self.yaw = yaw;
        self.pitch = pitch;

        self.update_restrictions();
        self.update_projviews();
    }

    /// Transformation applied by the camera without perspective.
    fn update_restrictions(&mut self) {
        if self.dist < self.min_dist {
            self.dist = self.min_dist
        }

        if self.dist > self.max_dist {
            self.dist = self.max_dist
        }

        if self.pitch <= self.min_pitch {
            self.pitch = self.min_pitch
        }

        if self.pitch > self.max_pitch {
            self.pitch = self.max_pitch
        }
    }

    /// The button used to rotate the ArcBall camera.
    pub fn rotate_button(&self) -> Option<MouseButton> {
        self.rotate_button
    }

    /// Set the button used to rotate the ArcBall camera.
    /// Use None to disable rotation.
    pub fn rebind_rotate_button(&mut self, new_button: Option<MouseButton>) {
        self.rotate_button = new_button;
    }

    /// Modifiers that must be pressed for the camera rotation to occur.
    pub fn rotate_modifiers(&self) -> Option<Modifiers> {
        self.rotate_modifiers
    }

    /// Sets the modifiers that must be pressed for the camera rotation to occur.
    ///
    /// If this is set to `None`, then pressing any modifier will not prevent rotation from occurring.
    /// If this is different from `None` then rotation will occur only if the exact specified set of modifiers is pressed.
    /// In particular, if this is set to `Some(Modifiers::empty())` then, rotation will occur only of no modifier is pressed.
    pub fn set_rotate_modifiers(&mut self, modifiers: Option<Modifiers>) {
        self.rotate_modifiers = modifiers
    }

    /// Modifiers that must be pressed for the camera drag to occur.
    pub fn drag_modifiers(&self) -> Option<Modifiers> {
        self.drag_modifiers
    }

    /// Sets the modifiers that must be pressed for the camera drag to occur.
    ///
    /// If this is set to `None`, then pressing any modifier will not prevent dragging from occurring.
    /// If this is different from `None` then drag will occur only if the exact specified set of modifiers is pressed.
    /// In particular, if this is set to `Some(Modifiers::empty())` then, drag will occur only of no modifier is pressed.
    pub fn set_drag_modifiers(&mut self, modifiers: Option<Modifiers>) {
        self.drag_modifiers = modifiers
    }

    /// The button used to drag the ArcBall camera.
    pub fn drag_button(&self) -> Option<MouseButton> {
        self.drag_button
    }

    /// Set the button used to drag the ArcBall camera.
    /// Use None to disable dragging.
    pub fn rebind_drag_button(&mut self, new_button: Option<MouseButton>) {
        self.drag_button = new_button;
    }

    /// The key used to reset the ArcBall camera.
    pub fn reset_key(&self) -> Option<Key> {
        self.reset_key
    }

    /// Set the key used to reset the ArcBall camera.
    /// Use None to disable reset.
    pub fn rebind_reset_key(&mut self, new_key: Option<Key>) {
        self.reset_key = new_key;
    }

    fn handle_left_button_displacement(&mut self, dpos: &Vector2<f32>) {
        self.yaw = self.yaw + dpos.x * self.yaw_step;
        self.pitch = self.pitch - dpos.y * self.pitch_step;

        self.update_restrictions();
        self.update_projviews();
    }

    fn handle_right_button_displacement(&mut self, dpos: &Vector2<f32>) {
        let eye = self.eye();
        let dir = (self.at - eye).normalize();
        let tangent = Vector3::y().cross(&dir).normalize();
        let bitangent = dir.cross(&tangent);
        let mult = self.dist / 1000.0;

        self.at = self.at + tangent * (dpos.x * mult) + bitangent * (dpos.y * mult);
        self.update_projviews();
    }

    fn handle_scroll(&mut self, off: f32) {
        self.dist = self.dist + self.dist_step * (off) / 120.0;
        self.update_restrictions();
        self.update_projviews();
    }

    fn update_projviews(&mut self) {
        self.proj = *self.projection.as_matrix();
        self.view = self.view_transform().to_homogeneous();
        self.proj_view = self.proj * self.view;
        self.inverse_proj_view = self.proj_view.try_inverse().unwrap();
    }

    /// Sets the up vector of this camera.
    #[inline]
    pub fn set_up_axis(&mut self, up_axis: Vector3<f32>) {
        self.up_axis = up_axis;
    }
}

impl Camera for ArcBall {
    fn clip_planes(&self) -> (f32, f32) {
        (self.projection.znear(), self.projection.zfar())
    }

    fn view_transform(&self) -> Isometry3<f32> {
        Isometry3::look_at_rh(&self.eye(), &self.at, &self.up_axis)
    }

    fn eye(&self) -> Point3<f32> {
        let px = self.at.x + self.dist * self.yaw.cos() * self.pitch.sin();
        let py = self.at.y + self.dist * self.pitch.cos();
        let pz = self.at.z + self.dist * self.yaw.sin() * self.pitch.sin();

        Point3::new(px, py, pz)
    }

    fn handle_event(&mut self, canvas: &Canvas, event: &WindowEvent) {
        match *event {
            WindowEvent::CursorPos(x, y, modifiers) => {
                let curr_pos = Vector2::new(x as f32, y as f32);

                if let Some(rotate_button) = self.rotate_button {
                    if canvas.get_mouse_button(rotate_button) == Action::Press &&
                        self.rotate_modifiers.map(|m| m == modifiers).unwrap_or(true) {
                        let dpos = curr_pos - self.last_cursor_pos;
                        self.handle_left_button_displacement(&dpos)
                    }
                }

                if let Some(drag_button) = self.drag_button {
                    if canvas.get_mouse_button(drag_button) == Action::Press &&
                        self.drag_modifiers.map(|m| m == modifiers).unwrap_or(true) {
                        let dpos = curr_pos - self.last_cursor_pos;
                        self.handle_right_button_displacement(&dpos)
                    }
                }

                self.last_cursor_pos = curr_pos;
            }
            WindowEvent::Key(key, Action::Press, _) if Some(key) == self.reset_key => {
                self.at = Point3::origin();
                self.update_projviews();
            }
            WindowEvent::Scroll(_, off, _) => self.handle_scroll(off as f32),
            WindowEvent::FramebufferSize(w, h) => {
                self.projection.set_aspect(w as f32 / h as f32);
                self.update_projviews();
            }
            _ => {}
        }
    }

    #[inline]
    fn upload(
        &self,
        _: usize,
        proj: &mut ShaderUniform<Matrix4<f32>>,
        view: &mut ShaderUniform<Matrix4<f32>>,
    ) {
        proj.upload(&self.proj);
        view.upload(&self.view);
    }

    fn transformation(&self) -> Matrix4<f32> {
        self.proj_view
    }

    fn inverse_transformation(&self) -> Matrix4<f32> {
        self.inverse_proj_view
    }

    fn update(&mut self, _: &Canvas) {}
}
<|endoftext|>


<commit_before>// Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

use protoc_grpcio;

use std::path::{Path, PathBuf};

use build_utils::BuildRoot;
use std::collections::HashSet;

fn main() {
  let build_root = BuildRoot::find().unwrap();
  let thirdpartyprotobuf = build_root.join("3rdparty/protobuf");
  mark_dir_as_rerun_trigger(&thirdpartyprotobuf);

  let grpcio_output_dir = PathBuf::from("src/gen");
  replace_if_changed(&grpcio_output_dir, |path| {
    generate_for_grpcio(&thirdpartyprotobuf, path);
    format(path);
  });

  let tower_output_dir = PathBuf::from("src/gen_for_tower");
  replace_if_changed(&tower_output_dir, |path| {
    generate_for_tower(&thirdpartyprotobuf, path);
    format(path);
  });

  // Re-gen if, say, someone does a git clean on the gen dir but not the target dir. This ensures
  // generated sources are available for reading by programmers and tools like rustfmt alike.
  mark_dir_as_rerun_trigger(&grpcio_output_dir);
  mark_dir_as_rerun_trigger(&tower_output_dir);
}

fn generate_for_grpcio(thirdpartyprotobuf: &Path, gen_dir: &Path) {
  let amended_proto_root =
    add_rustproto_header(&thirdpartyprotobuf).expect("Error adding proto bytes header");

  protoc_grpcio::compile_grpc_protos(
    &[
      "build/bazel/remote/execution/v2/remote_execution.proto",
      "google/bytestream/bytestream.proto",
      "google/rpc/code.proto",
      "google/rpc/error_details.proto",
      "google/rpc/status.proto",
      "google/longrunning/operations.proto",
      "google/protobuf/empty.proto",
    ],
    &[
      amended_proto_root.path().to_owned(),
      thirdpartyprotobuf.join("standard"),
      thirdpartyprotobuf.join("rust-protobuf"),
    ],
    &gen_dir,
  )
  .expect("Failed to compile protos!");

  disable_clippy_in_generated_code(&gen_dir).expect("Failed to strip clippy from generated code");

  generate_mod_rs(&gen_dir).expect("Failed to generate mod.rs");
}

fn mark_dir_as_rerun_trigger(dir: &Path) {
  for file in walkdir::WalkDir::new(dir) {
    println!("cargo:rerun-if-changed={}", file.unwrap().path().display());
  }
}

const EXTRA_HEADER: &'static str = r#"import "rustproto.proto";
option (rustproto.carllerche_bytes_for_bytes_all) = true;
"#;

///
/// Copies protos from thirdpartyprotobuf, adds a header to make protoc_grpcio uses Bytes instead
/// of Vec<u8>s, and rewrites them into a temporary directory
///
fn add_rustproto_header(thirdpartyprotobuf: &Path) -> Result<tempfile::TempDir, String> {
  let amended_proto_root = tempfile::TempDir::new().unwrap();
  for f in &["bazelbuild_remote-apis", "googleapis"] {
    let src_root = thirdpartyprotobuf.join(f);
    for entry in walkdir::WalkDir::new(&src_root)
      .into_iter()
      .filter_map(|entry| entry.ok())
      .filter(|entry| entry.file_type().is_file())
      .filter(|entry| entry.file_name().to_string_lossy().ends_with(".proto"))
    {
      let dst = amended_proto_root
        .path()
        .join(entry.path().strip_prefix(&src_root).unwrap());
      std::fs::create_dir_all(dst.parent().unwrap())
        .map_err(|err| format!("Error making dir in temp proto root: {}", err))?;
      let original = std::fs::read_to_string(entry.path())
        .map_err(|err| format!("Error reading proto {}: {}", entry.path().display(), err))?;
      let mut copy = String::with_capacity(original.len() + EXTRA_HEADER.len());
      for line in original.lines() {
        copy += line;
        copy += "\n";
        if line.starts_with("package ") {
          copy += EXTRA_HEADER
        }
      }
      std::fs::write(&dst, copy.as_bytes())
        .map_err(|err| format!("Error writing {}: {}", dst.display(), err))?;
    }
  }
  Ok(amended_proto_root)
}

///
/// protoc_grpcio generates its own clippy config, but it's for an out of date version of clippy,
/// so strip that out so we don't get warnings about it.
///
/// Add our own #![allow(clippy::all)] heading to each generated file so that we don't get any
/// warnings/errors from generated code not meeting our standards.
///
fn disable_clippy_in_generated_code(dir: &Path) -> Result<(), String> {
  for file in walkdir::WalkDir::new(&dir)
    .into_iter()
    .filter_map(|entry| entry.ok())
    .filter(|entry| {
      entry.file_type().is_file() && entry.file_name().to_string_lossy().ends_with(".rs")
    })
  {
    let lines: Vec<_> = std::fs::read_to_string(file.path())
      .map_err(|err| {
        format!(
          "Error reading generated protobuf at {}: {}",
          file.path().display(),
          err
        )
      })?
      .lines()
      .filter(|line| !line.contains("clippy"))
      .map(str::to_owned)
      .collect();
    let content = String::from("#![allow(clippy::all)]\n") + &lines.join("\n");
    std::fs::write(file.path(), content).map_err(|err| {
      format!(
        "Error re-writing generated protobuf at {}: {}",
        file.path().display(),
        err
      )
    })?;
  }
  Ok(())
}

fn generate_mod_rs(dir: &Path) -> Result<(), String> {
  let listing = dir.read_dir().unwrap();
  let mut pub_mod_stmts = listing
    .filter_map(|d| d.ok())
    .map(|d| d.file_name().to_string_lossy().into_owned())
    .filter(|name| &name != &"mod.rs" && &name != &".gitignore")
    .map(|name| format!("pub mod {};", name.trim_end_matches(".rs")))
    .collect::<Vec<_>>();
  pub_mod_stmts.sort();
  let contents = format!(
    "\
// This file is generated. Do not edit.
{}
",
    pub_mod_stmts.join("\n")
  );

  std::fs::write(dir.join("mod.rs"), contents)
    .map_err(|err| format!("Failed to write mod.rs: {}", err))
}

fn generate_for_tower(thirdpartyprotobuf: &Path, out_dir: &Path) {
  tower_grpc_build::Config::new()
    .enable_server(true)
    .enable_client(true)
    .build(
      &[PathBuf::from(
        "build/bazel/remote/execution/v2/remote_execution.proto",
      )],
      &std::fs::read_dir(&thirdpartyprotobuf)
        .unwrap()
        .into_iter()
        .map(|d| d.unwrap().path())
        .collect::<Vec<_>>(),
    )
    .unwrap_or_else(|e| panic!("protobuf compilation failed: {}", e));

  let mut dirs_needing_mod_rs = HashSet::new();
  dirs_needing_mod_rs.insert(out_dir.to_owned());

  for f in walkdir::WalkDir::new(std::env::var("OUT_DIR").unwrap())
    .into_iter()
    .filter_map(|f| f.ok())
    .filter(|f| f.path().extension() == Some("rs".as_ref()))
  {
    let mut parts: Vec<_> = f
      .path()
      .file_name()
      .unwrap()
      .to_str()
      .unwrap()
      .split('.')
      .collect();
    // pop .rs
    parts.pop();

    let mut dst = out_dir.to_owned();
    for part in parts {
      dst.push(part);
      dirs_needing_mod_rs.insert(dst.clone());
      if !dst.exists() {
        std::fs::create_dir_all(&dst).unwrap();
      }
    }
    dirs_needing_mod_rs.remove(&dst);
    dst = dst.join("mod.rs");

    std::fs::copy(f.path(), dst).unwrap();
  }

  disable_clippy_in_generated_code(out_dir).expect("Failed to strip clippy from generated code");

  for dir in &dirs_needing_mod_rs {
    generate_mod_rs(dir).expect("Failed to write mod.rs");
  }
}

///
/// Replaces contents of directory `path` with contents of directory populated by passed function.
/// Does not modify `path` if the contents are identical.
///
fn replace_if_changed<F: FnOnce(&Path)>(path: &Path, f: F) {
  let tempdir = tempfile::TempDir::new().unwrap();
  f(tempdir.path());
  if !dir_diff::is_different(path, tempdir.path()).unwrap() {
    return;
  }
  if path.exists() {
    std::fs::remove_dir_all(path).unwrap();
  }
  std::fs::create_dir_all(path.parent().unwrap()).unwrap();
  copy_dir::copy_dir(tempdir.path(), path).unwrap();
}

fn format(path: &Path) {
  let mut rustfmt = PathBuf::from(env!("CARGO"));
  rustfmt.pop();
  rustfmt.push("rustfmt");
  let mut rustfmt_config = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
  rustfmt_config.pop(); // bazel_protos
  rustfmt_config.pop(); // process_execution
  rustfmt_config.push("rustfmt.toml");

  if !rustfmt_config.exists() {
    panic!("Couldn't find file {}", rustfmt_config.display());
  }

  for file in walkdir::WalkDir::new(path) {
    let file = file.unwrap();
    if file.file_type().is_file() {
      let success = std::process::Command::new(&rustfmt)
        .arg(file.path())
        .arg("--config-path")
        .arg(&rustfmt_config)
        .status()
        .unwrap()
        .success();
      if !success {
        panic!("Cargo formatting failed for generated protos. Output should be above.");
      }
    }
  }
}
<commit_msg>Quash warnings about missing dyn in generated code (#8245)<commit_after>// Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

use protoc_grpcio;

use std::path::{Path, PathBuf};

use build_utils::BuildRoot;
use std::collections::HashSet;

fn main() {
  let build_root = BuildRoot::find().unwrap();
  let thirdpartyprotobuf = build_root.join("3rdparty/protobuf");
  mark_dir_as_rerun_trigger(&thirdpartyprotobuf);

  let grpcio_output_dir = PathBuf::from("src/gen");
  replace_if_changed(&grpcio_output_dir, |path| {
    generate_for_grpcio(&thirdpartyprotobuf, path);
    format(path);
  });

  let tower_output_dir = PathBuf::from("src/gen_for_tower");
  replace_if_changed(&tower_output_dir, |path| {
    generate_for_tower(&thirdpartyprotobuf, path);
    format(path);
  });

  // Re-gen if, say, someone does a git clean on the gen dir but not the target dir. This ensures
  // generated sources are available for reading by programmers and tools like rustfmt alike.
  mark_dir_as_rerun_trigger(&grpcio_output_dir);
  mark_dir_as_rerun_trigger(&tower_output_dir);
}

fn generate_for_grpcio(thirdpartyprotobuf: &Path, gen_dir: &Path) {
  let amended_proto_root =
    add_rustproto_header(&thirdpartyprotobuf).expect("Error adding proto bytes header");

  protoc_grpcio::compile_grpc_protos(
    &[
      "build/bazel/remote/execution/v2/remote_execution.proto",
      "google/bytestream/bytestream.proto",
      "google/rpc/code.proto",
      "google/rpc/error_details.proto",
      "google/rpc/status.proto",
      "google/longrunning/operations.proto",
      "google/protobuf/empty.proto",
    ],
    &[
      amended_proto_root.path().to_owned(),
      thirdpartyprotobuf.join("standard"),
      thirdpartyprotobuf.join("rust-protobuf"),
    ],
    &gen_dir,
  )
  .expect("Failed to compile protos!");

  disable_clippy_in_generated_code(&gen_dir).expect("Failed to strip clippy from generated code");

  generate_mod_rs(&gen_dir).expect("Failed to generate mod.rs");
}

fn mark_dir_as_rerun_trigger(dir: &Path) {
  for file in walkdir::WalkDir::new(dir) {
    println!("cargo:rerun-if-changed={}", file.unwrap().path().display());
  }
}

const EXTRA_HEADER: &'static str = r#"import "rustproto.proto";
option (rustproto.carllerche_bytes_for_bytes_all) = true;
"#;

///
/// Copies protos from thirdpartyprotobuf, adds a header to make protoc_grpcio uses Bytes instead
/// of Vec<u8>s, and rewrites them into a temporary directory
///
fn add_rustproto_header(thirdpartyprotobuf: &Path) -> Result<tempfile::TempDir, String> {
  let amended_proto_root = tempfile::TempDir::new().unwrap();
  for f in &["bazelbuild_remote-apis", "googleapis"] {
    let src_root = thirdpartyprotobuf.join(f);
    for entry in walkdir::WalkDir::new(&src_root)
      .into_iter()
      .filter_map(|entry| entry.ok())
      .filter(|entry| entry.file_type().is_file())
      .filter(|entry| entry.file_name().to_string_lossy().ends_with(".proto"))
    {
      let dst = amended_proto_root
        .path()
        .join(entry.path().strip_prefix(&src_root).unwrap());
      std::fs::create_dir_all(dst.parent().unwrap())
        .map_err(|err| format!("Error making dir in temp proto root: {}", err))?;
      let original = std::fs::read_to_string(entry.path())
        .map_err(|err| format!("Error reading proto {}: {}", entry.path().display(), err))?;
      let mut copy = String::with_capacity(original.len() + EXTRA_HEADER.len());
      for line in original.lines() {
        copy += line;
        copy += "\n";
        if line.starts_with("package ") {
          copy += EXTRA_HEADER
        }
      }
      std::fs::write(&dst, copy.as_bytes())
        .map_err(|err| format!("Error writing {}: {}", dst.display(), err))?;
    }
  }
  Ok(amended_proto_root)
}

///
/// protoc_grpcio generates its own clippy config, but it's for an out of date version of clippy,
/// so strip that out so we don't get warnings about it.
///
/// Add our own #![allow(clippy::all)] heading to each generated file so that we don't get any
/// warnings/errors from generated code not meeting our standards.
///
fn disable_clippy_in_generated_code(dir: &Path) -> Result<(), String> {
  for file in walkdir::WalkDir::new(&dir)
    .into_iter()
    .filter_map(|entry| entry.ok())
    .filter(|entry| {
      entry.file_type().is_file() && entry.file_name().to_string_lossy().ends_with(".rs")
    })
  {
    let lines: Vec<_> = std::fs::read_to_string(file.path())
      .map_err(|err| {
        format!(
          "Error reading generated protobuf at {}: {}",
          file.path().display(),
          err
        )
      })?
      .lines()
      .filter(|line| !line.contains("clippy"))
      .map(str::to_owned)
      .collect();
    let content = (String::from("#![allow(clippy::all)]\n") + &lines.join("\n"))
      // Quash warnings about missing dyn, because we can't currently update the code generator
      // to get rid of them because we forked it upstream a while ago.
      .replace("::std::any::Any", "dyn ::std::any::Any");
    std::fs::write(file.path(), content).map_err(|err| {
      format!(
        "Error re-writing generated protobuf at {}: {}",
        file.path().display(),
        err
      )
    })?;
  }
  Ok(())
}

fn generate_mod_rs(dir: &Path) -> Result<(), String> {
  let listing = dir.read_dir().unwrap();
  let mut pub_mod_stmts = listing
    .filter_map(|d| d.ok())
    .map(|d| d.file_name().to_string_lossy().into_owned())
    .filter(|name| &name != &"mod.rs" && &name != &".gitignore")
    .map(|name| format!("pub mod {};", name.trim_end_matches(".rs")))
    .collect::<Vec<_>>();
  pub_mod_stmts.sort();
  let contents = format!(
    "\
// This file is generated. Do not edit.
{}
",
    pub_mod_stmts.join("\n")
  );

  std::fs::write(dir.join("mod.rs"), contents)
    .map_err(|err| format!("Failed to write mod.rs: {}", err))
}

fn generate_for_tower(thirdpartyprotobuf: &Path, out_dir: &Path) {
  tower_grpc_build::Config::new()
    .enable_server(true)
    .enable_client(true)
    .build(
      &[PathBuf::from(
        "build/bazel/remote/execution/v2/remote_execution.proto",
      )],
      &std::fs::read_dir(&thirdpartyprotobuf)
        .unwrap()
        .into_iter()
        .map(|d| d.unwrap().path())
        .collect::<Vec<_>>(),
    )
    .unwrap_or_else(|e| panic!("protobuf compilation failed: {}", e));

  let mut dirs_needing_mod_rs = HashSet::new();
  dirs_needing_mod_rs.insert(out_dir.to_owned());

  for f in walkdir::WalkDir::new(std::env::var("OUT_DIR").unwrap())
    .into_iter()
    .filter_map(|f| f.ok())
    .filter(|f| f.path().extension() == Some("rs".as_ref()))
  {
    let mut parts: Vec<_> = f
      .path()
      .file_name()
      .unwrap()
      .to_str()
      .unwrap()
      .split('.')
      .collect();
    // pop .rs
    parts.pop();

    let mut dst = out_dir.to_owned();
    for part in parts {
      dst.push(part);
      dirs_needing_mod_rs.insert(dst.clone());
      if !dst.exists() {
        std::fs::create_dir_all(&dst).unwrap();
      }
    }
    dirs_needing_mod_rs.remove(&dst);
    dst = dst.join("mod.rs");

    std::fs::copy(f.path(), dst).unwrap();
  }

  disable_clippy_in_generated_code(out_dir).expect("Failed to strip clippy from generated code");

  for dir in &dirs_needing_mod_rs {
    generate_mod_rs(dir).expect("Failed to write mod.rs");
  }
}

///
/// Replaces contents of directory `path` with contents of directory populated by passed function.
/// Does not modify `path` if the contents are identical.
///
fn replace_if_changed<F: FnOnce(&Path)>(path: &Path, f: F) {
  let tempdir = tempfile::TempDir::new().unwrap();
  f(tempdir.path());
  if !dir_diff::is_different(path, tempdir.path()).unwrap() {
    return;
  }
  if path.exists() {
    std::fs::remove_dir_all(path).unwrap();
  }
  std::fs::create_dir_all(path.parent().unwrap()).unwrap();
  copy_dir::copy_dir(tempdir.path(), path).unwrap();
}

fn format(path: &Path) {
  let mut rustfmt = PathBuf::from(env!("CARGO"));
  rustfmt.pop();
  rustfmt.push("rustfmt");
  let mut rustfmt_config = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
  rustfmt_config.pop(); // bazel_protos
  rustfmt_config.pop(); // process_execution
  rustfmt_config.push("rustfmt.toml");

  if !rustfmt_config.exists() {
    panic!("Couldn't find file {}", rustfmt_config.display());
  }

  for file in walkdir::WalkDir::new(path) {
    let file = file.unwrap();
    if file.file_type().is_file() {
      let success = std::process::Command::new(&rustfmt)
        .arg(file.path())
        .arg("--config-path")
        .arg(&rustfmt_config)
        .status()
        .unwrap()
        .success();
      if !success {
        panic!("Cargo formatting failed for generated protos. Output should be above.");
      }
    }
  }
}
<|endoftext|>


<commit_before>//! Encoding of AVIF images.
///
/// The [AVIF] specification defines an image derivative of the AV1 bitstream, an open video codec.
///
/// [AVIF]: https://aomediacodec.github.io/av1-avif/
use std::borrow::Cow;
use std::io::Write;

use crate::{ColorType, ImageBuffer, ImageFormat, Pixel};
use crate::{ImageError, ImageResult};
use crate::buffer::ConvertBuffer;
use crate::color::{FromColor, Luma, LumaA, Bgr, Bgra, Rgb, Rgba};
use crate::error::{EncodingError, ParameterError, ParameterErrorKind, UnsupportedError, UnsupportedErrorKind};

use bytemuck::{Pod, PodCastError, try_cast_slice, try_cast_slice_mut};
use num_traits::Zero;
use ravif::{Img, ColorSpace, Config, RGBA8, encode_rgba};
use rgb::AsPixels;

/// AVIF Encoder.
///
/// Writes one image into the chosen output.
pub struct AvifEncoder<W> {
    inner: W,
    fallback: Vec<u8>,
}

impl<W: Write> AvifEncoder<W> {
    /// Create a new encoder that writes its output to `w`.
    pub fn new(w: W) -> Self {
        AvifEncoder { inner: w, fallback: vec![] }
    }

    /// Encode image data with the indicated color type.
    ///
    /// The encoder currently requires all data to be RGBA8, it will be converted internally if
    /// necessary. When data is suitably aligned, i.e. u16 channels to two bytes, then the
    /// conversion may be more efficient.
    pub fn write_image(mut self, data: &[u8], width: u32, height: u32, color: ColorType) -> ImageResult<()> {
        let config = self.config(color);
        // `ravif` needs strongly typed data so let's convert. We can either use a temporarily
        // owned version in our own buffer or zero-copy if possible by using the input buffer.
        // This requires going through `rgb`.
        let buffer = self.encode_as_img(data, width, height, color)?;
        let (data, _color_size, _alpha_size) = encode_rgba(buffer, &config)
            .map_err(|err| ImageError::Encoding(
                EncodingError::new(ImageFormat::Avif.into(), err)
            ))?;
        self.inner.write_all(&data)?;
        Ok(())
    }

    fn config(&self, _color: ColorType) -> Config {
        Config {
            quality: 100,
            alpha_quality: 100,
            speed: 1,
            premultiplied_alpha: false,
            color_space: ColorSpace::RGB,
        }
    }

    fn encode_as_img<'buf>(&'buf mut self, data: &'buf [u8], width: u32, height: u32, color: ColorType)
        -> ImageResult<Img<&'buf [RGBA8]>>
    {
        // Error wrapping utility for color dependent buffer dimensions.
        fn try_from_raw<P: Pixel + 'static>(data: &[P::Subpixel], width: u32, height: u32)
            -> ImageResult<ImageBuffer<P, &[P::Subpixel]>>
        {
            ImageBuffer::from_raw(width, height, data).ok_or_else(|| {
                ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch))
            })
        };

        // Convert to target color type using few buffer allocations.
        fn convert_into<'buf, P>(buf: &'buf mut Vec<u8>, image: ImageBuffer<P, &[P::Subpixel]>)
            -> Img<&'buf [RGBA8]>
        where
            P: Pixel + 'static,
            Rgba<u8>: FromColor<P>,
        {
            let (width, height) = image.dimensions();
            // TODO: conversion re-using the target buffer?
            let image: ImageBuffer<Rgba<u8>, _> = image.convert();
            *buf = image.into_raw();
            Img::new(buf.as_pixels(), width as usize, height as usize)
        }

        // Cast the input slice using few buffer allocations if possible.
        // In particular try not to allocate if the caller did the infallible reverse.
        fn cast_buffer<Channel>(buf: &[u8]) -> ImageResult<Cow<[Channel]>>
        where
            Channel: Pod + Zero,
        {
            match try_cast_slice(buf) {
                Ok(slice) => Ok(Cow::Borrowed(slice)),
                Err(PodCastError::OutputSliceWouldHaveSlop) => {
                    Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                }
                Err(PodCastError::TargetAlignmentGreaterAndInputNotAligned) => {
                    // Sad, but let's allocate.
                    // bytemuck checks alignment _before_ slop but size mismatch before this..
                    if buf.len() % std::mem::size_of::<Channel>() != 0 {
                        Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                    } else {
                        let len = buf.len() / std::mem::size_of::<Channel>();
                        let mut data = vec![Channel::zero(); len];
                        let view = try_cast_slice_mut::<_, u8>(data.as_mut_slice()).unwrap();
                        view.copy_from_slice(buf);
                        Ok(Cow::Owned(data))
                    }
                }
                Err(err) => { // Are you trying to encode a ZST??
                    Err(ImageError::Parameter(ParameterError::from_kind(
                        ParameterErrorKind::Generic(format!("{:?}", err))
                    )))
                }
            }
        }

        match color {
            ColorType::Rgba8 => {
                // ravif doesn't do any checks but has some asserts, so we do the checks.
                let img = try_from_raw::<Rgba<u8>>(data, width, height)?;
                // Now, internally ravif uses u32 but it takes usize. We could do some checked
                // conversion but instead we use that a non-empty image must be addressable.
                if img.pixels().len() == 0 {
                    return Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)));
                }

                Ok(Img::new(rgb::AsPixels::as_pixels(data), width as usize, height as usize))
            },
            // we need a separate buffer..
            ColorType::L8 => {
                let image = try_from_raw::<Luma<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::La8 => {
                let image = try_from_raw::<LumaA<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgb8 => {
                let image = try_from_raw::<Rgb<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Bgr8 => {
                let image = try_from_raw::<Bgr<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Bgra8 => {
                let image = try_from_raw::<Bgra<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            // we need to really convert data..
            ColorType::L16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Luma<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::La16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<LumaA<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgb16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgb<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgba16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgba<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            // for cases we do not support at all?
            _ => Err(ImageError::Unsupported(UnsupportedError::from_format_and_kind(
                    ImageFormat::Avif.into(),
                    UnsupportedErrorKind::Color(color.into()),
                )))
        }
    }
}
<commit_msg>Add write_image_with_speed_quality function.<commit_after>//! Encoding of AVIF images.
///
/// The [AVIF] specification defines an image derivative of the AV1 bitstream, an open video codec.
///
/// [AVIF]: https://aomediacodec.github.io/av1-avif/
use std::borrow::Cow;
use std::io::Write;
use std::cmp::{min, max};

use crate::{ColorType, ImageBuffer, ImageFormat, Pixel};
use crate::{ImageError, ImageResult};
use crate::buffer::ConvertBuffer;
use crate::color::{FromColor, Luma, LumaA, Bgr, Bgra, Rgb, Rgba};
use crate::error::{EncodingError, ParameterError, ParameterErrorKind, UnsupportedError, UnsupportedErrorKind};

use bytemuck::{Pod, PodCastError, try_cast_slice, try_cast_slice_mut};
use num_traits::Zero;
use ravif::{Img, ColorSpace, Config, RGBA8, encode_rgba};
use rgb::AsPixels;

/// AVIF Encoder.
///
/// Writes one image into the chosen output.
pub struct AvifEncoder<W> {
    inner: W,
    fallback: Vec<u8>,
}

impl<W: Write> AvifEncoder<W> {
    /// Create a new encoder that writes its output to `w`.
    pub fn new(w: W) -> Self {
        AvifEncoder { inner: w, fallback: vec![] }
    }

    /// Encode image data with the indicated color type.
    ///
    /// The encoder currently requires all data to be RGBA8, it will be converted internally if
    /// necessary. When data is suitably aligned, i.e. u16 channels to two bytes, then the
    /// conversion may be more efficient.
    pub fn write_image(mut self, data: &[u8], width: u32, height: u32, color: ColorType) -> ImageResult<()> {
        self.write_image_with_speed_quality(data, width, height, color, 100, 1)
    }

    /// Encode image data with the indicated color type, quality and speed.
    ///
    /// The encoder currently requires all data to be RGBA8, it will be converted internally if
    /// necessary. When data is suitably aligned, i.e. u16 channels to two bytes, then the
    /// conversion may be more efficient.
    pub fn write_image_with_speed_quality(mut self, data: &[u8], width: u32, height: u32, color: ColorType, quality: u8, speed: u8) -> ImageResult<()> {
        // Clamp quality and speed to range
        let quality = max(1, min(quality, 100));
        let speed = max(1, min(speed, 10));

        let config = self.config(color, quality, speed);
        // `ravif` needs strongly typed data so let's convert. We can either use a temporarily
        // owned version in our own buffer or zero-copy if possible by using the input buffer.
        // This requires going through `rgb`.
        let buffer = self.encode_as_img(data, width, height, color)?;
        let (data, _color_size, _alpha_size) = encode_rgba(buffer, &config)
            .map_err(|err| ImageError::Encoding(
                EncodingError::new(ImageFormat::Avif.into(), err)
            ))?;
        self.inner.write_all(&data)?;
        Ok(())
    }

    fn config(&self, _color: ColorType, quality: u8, speed: u8) -> Config {
        Config {
            quality: quality,
            alpha_quality: quality,
            speed: speed,
            premultiplied_alpha: false,
            color_space: ColorSpace::RGB,
        }
    }

    fn encode_as_img<'buf>(&'buf mut self, data: &'buf [u8], width: u32, height: u32, color: ColorType)
        -> ImageResult<Img<&'buf [RGBA8]>>
    {
        // Error wrapping utility for color dependent buffer dimensions.
        fn try_from_raw<P: Pixel + 'static>(data: &[P::Subpixel], width: u32, height: u32)
            -> ImageResult<ImageBuffer<P, &[P::Subpixel]>>
        {
            ImageBuffer::from_raw(width, height, data).ok_or_else(|| {
                ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch))
            })
        };

        // Convert to target color type using few buffer allocations.
        fn convert_into<'buf, P>(buf: &'buf mut Vec<u8>, image: ImageBuffer<P, &[P::Subpixel]>)
            -> Img<&'buf [RGBA8]>
        where
            P: Pixel + 'static,
            Rgba<u8>: FromColor<P>,
        {
            let (width, height) = image.dimensions();
            // TODO: conversion re-using the target buffer?
            let image: ImageBuffer<Rgba<u8>, _> = image.convert();
            *buf = image.into_raw();
            Img::new(buf.as_pixels(), width as usize, height as usize)
        }

        // Cast the input slice using few buffer allocations if possible.
        // In particular try not to allocate if the caller did the infallible reverse.
        fn cast_buffer<Channel>(buf: &[u8]) -> ImageResult<Cow<[Channel]>>
        where
            Channel: Pod + Zero,
        {
            match try_cast_slice(buf) {
                Ok(slice) => Ok(Cow::Borrowed(slice)),
                Err(PodCastError::OutputSliceWouldHaveSlop) => {
                    Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                }
                Err(PodCastError::TargetAlignmentGreaterAndInputNotAligned) => {
                    // Sad, but let's allocate.
                    // bytemuck checks alignment _before_ slop but size mismatch before this..
                    if buf.len() % std::mem::size_of::<Channel>() != 0 {
                        Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                    } else {
                        let len = buf.len() / std::mem::size_of::<Channel>();
                        let mut data = vec![Channel::zero(); len];
                        let view = try_cast_slice_mut::<_, u8>(data.as_mut_slice()).unwrap();
                        view.copy_from_slice(buf);
                        Ok(Cow::Owned(data))
                    }
                }
                Err(err) => { // Are you trying to encode a ZST??
                    Err(ImageError::Parameter(ParameterError::from_kind(
                        ParameterErrorKind::Generic(format!("{:?}", err))
                    )))
                }
            }
        }

        match color {
            ColorType::Rgba8 => {
                // ravif doesn't do any checks but has some asserts, so we do the checks.
                let img = try_from_raw::<Rgba<u8>>(data, width, height)?;
                // Now, internally ravif uses u32 but it takes usize. We could do some checked
                // conversion but instead we use that a non-empty image must be addressable.
                if img.pixels().len() == 0 {
                    return Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)));
                }

                Ok(Img::new(rgb::AsPixels::as_pixels(data), width as usize, height as usize))
            },
            // we need a separate buffer..
            ColorType::L8 => {
                let image = try_from_raw::<Luma<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::La8 => {
                let image = try_from_raw::<LumaA<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgb8 => {
                let image = try_from_raw::<Rgb<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Bgr8 => {
                let image = try_from_raw::<Bgr<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Bgra8 => {
                let image = try_from_raw::<Bgra<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            // we need to really convert data..
            ColorType::L16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Luma<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::La16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<LumaA<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgb16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgb<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgba16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgba<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            // for cases we do not support at all?
            _ => Err(ImageError::Unsupported(UnsupportedError::from_format_and_kind(
                    ImageFormat::Avif.into(),
                    UnsupportedErrorKind::Color(color.into()),
                )))
        }
    }
}
<|endoftext|>