<commit_before>use super::command::{Command, Query};
use super::command::Command::LogicalShardList;
use std::collections::HashMap;
use command_query::CommandQuery;
use queryable::Queryable;
use command_line::CommandLine;
use commandable::Commandable;
use extendable::Extendable;
use request_cancellable::RequestCancellable;
use request_timeoutable::RequestTimeoutable;

#[derive(Clone, PartialEq, Eq, Debug)]
pub struct LogicalShardListCommand {
    command: Command,
    logical_table: String,
    arguments: HashMap<String, String>,
}

impl Default for LogicalShardListCommand {
    fn default() -> LogicalShardListCommand {
        LogicalShardListCommand {
            command: LogicalShardList,
            logical_table: "".to_string(),
            arguments: HashMap::new(),
        }
    }
}

impl LogicalShardListCommand {
    pub fn new(logical_table: String) -> LogicalShardListCommand {
        let default: LogicalShardListCommand = Default::default();
        LogicalShardListCommand {
            command: default.command,
            logical_table: logical_table,
            arguments: default.arguments,
        }
    }

    pub fn build(self) -> (Command, Query) {
        let mut query: Query = vec![("logical_table".to_string(), self.logical_table)];
        for (key, value) in &self.arguments {
            query.push((key.to_owned(), value.to_owned()));
        }
        (LogicalShardList, query)
    }
}

impl Queryable for LogicalShardListCommand {
    fn to_query(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandQuery::new(command, query);
        command.encode()
    }
}

impl Commandable for LogicalShardListCommand {
    fn to_command(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandLine::new(command, query);
        command.encode()
    }
}

extendable!(LogicalShardListCommand);
request_cancellable!(LogicalShardListCommand);
request_timeoutable!(LogicalShardListCommand);

#[cfg(test)]
mod test {
    use super::*;
    use command::Query;
    use command::Command::LogicalShardList;
    use std::collections::HashMap;
    use queryable::Queryable;
    use commandable::Commandable;
    use extendable::Extendable;

    #[test]
    fn test_new() {
        let table_list = LogicalShardListCommand::new("Logs".to_string());
        let expected = LogicalShardListCommand {
            command: LogicalShardList,
            logical_table: "Logs".to_string(),
            arguments: HashMap::new(),
        };
        assert_eq!(expected, table_list);
    }

    #[test]
    fn test_build() {
        let actual = LogicalShardListCommand::new("Logs".to_string()).build();
        let expected_query: Query = vec![("logical_table".to_string(), "Logs".to_string())];
        let expected = (LogicalShardList, expected_query);
        assert_eq!(expected, actual);
    }

    #[test]
    fn test_queryable() {
        let query = LogicalShardListCommand::new("Logs".to_string()).to_query();
        let url_encoded = "/d/logical_shard_list?logical_table=Logs";
        assert_eq!(url_encoded.to_string(), query);
    }

    #[test]
    fn test_commandable() {
        let query = LogicalShardListCommand::new("Logs".to_string()).to_command();
        let cli_encoded = "logical_shard_list --logical_table Logs";
        assert_eq!(cli_encoded.to_string(), query);
    }

    #[test]
    fn test_extendable() {
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("user".to_string(), "defined".to_string());
        let expected = LogicalShardListCommand {
            command: LogicalShardList,
            logical_table: "Logs".to_string(),
            arguments: arg.to_owned(),
        };
        let query = LogicalShardListCommand::new("Logs".to_string());
        unsafe {
            let extended = query.set_arguments(arg.to_owned());
            assert_eq!(expected, extended);
        }
    }
}
<commit_msg>Fix a typo<commit_after>use super::command::{Command, Query};
use super::command::Command::LogicalShardList;
use std::collections::HashMap;
use command_query::CommandQuery;
use queryable::Queryable;
use command_line::CommandLine;
use commandable::Commandable;
use extendable::Extendable;
use request_cancellable::RequestCancellable;
use request_timeoutable::RequestTimeoutable;

#[derive(Clone, PartialEq, Eq, Debug)]
pub struct LogicalShardListCommand {
    command: Command,
    logical_table: String,
    arguments: HashMap<String, String>,
}

impl Default for LogicalShardListCommand {
    fn default() -> LogicalShardListCommand {
        LogicalShardListCommand {
            command: LogicalShardList,
            logical_table: "".to_string(),
            arguments: HashMap::new(),
        }
    }
}

impl LogicalShardListCommand {
    pub fn new(logical_table: String) -> LogicalShardListCommand {
        let default: LogicalShardListCommand = Default::default();
        LogicalShardListCommand {
            command: default.command,
            logical_table: logical_table,
            arguments: default.arguments,
        }
    }

    pub fn build(self) -> (Command, Query) {
        let mut query: Query = vec![("logical_table".to_string(), self.logical_table)];
        for (key, value) in &self.arguments {
            query.push((key.to_owned(), value.to_owned()));
        }
        (LogicalShardList, query)
    }
}

impl Queryable for LogicalShardListCommand {
    fn to_query(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandQuery::new(command, query);
        command.encode()
    }
}

impl Commandable for LogicalShardListCommand {
    fn to_command(self) -> String {
        let (command, query) = self.build();
        let mut command = CommandLine::new(command, query);
        command.encode()
    }
}

extendable!(LogicalShardListCommand);
request_cancellable!(LogicalShardListCommand);
request_timeoutable!(LogicalShardListCommand);

#[cfg(test)]
mod test {
    use super::*;
    use command::Query;
    use command::Command::LogicalShardList;
    use std::collections::HashMap;
    use queryable::Queryable;
    use commandable::Commandable;
    use extendable::Extendable;

    #[test]
    fn test_new() {
        let logical_shard_list = LogicalShardListCommand::new("Logs".to_string());
        let expected = LogicalShardListCommand {
            command: LogicalShardList,
            logical_table: "Logs".to_string(),
            arguments: HashMap::new(),
        };
        assert_eq!(expected, logical_shard_list);
    }

    #[test]
    fn test_build() {
        let actual = LogicalShardListCommand::new("Logs".to_string()).build();
        let expected_query: Query = vec![("logical_table".to_string(), "Logs".to_string())];
        let expected = (LogicalShardList, expected_query);
        assert_eq!(expected, actual);
    }

    #[test]
    fn test_queryable() {
        let query = LogicalShardListCommand::new("Logs".to_string()).to_query();
        let url_encoded = "/d/logical_shard_list?logical_table=Logs";
        assert_eq!(url_encoded.to_string(), query);
    }

    #[test]
    fn test_commandable() {
        let query = LogicalShardListCommand::new("Logs".to_string()).to_command();
        let cli_encoded = "logical_shard_list --logical_table Logs";
        assert_eq!(cli_encoded.to_string(), query);
    }

    #[test]
    fn test_extendable() {
        let mut arg: HashMap<String, String> = HashMap::new();
        arg.insert("user".to_string(), "defined".to_string());
        let expected = LogicalShardListCommand {
            command: LogicalShardList,
            logical_table: "Logs".to_string(),
            arguments: arg.to_owned(),
        };
        let query = LogicalShardListCommand::new("Logs".to_string());
        unsafe {
            let extended = query.set_arguments(arg.to_owned());
            assert_eq!(expected, extended);
        }
    }
}
<|endoftext|>


<commit_before>use release_group::ReleaseGroup;
use uuid::Uuid;
use enums::{PersonType, AlbumMainType, AlbumSecondaryType};
use std::collections::HashMap;

#[derive(Debug, Clone)]
pub struct Artist {
    pub id: Uuid,
    pub name: String,
    pub gender: String,
    pub artist_type: PersonType,
    pub tags: Vec<String>,
    pub release_groups: Vec<ReleaseGroup>
}

impl Artist {
    pub fn new(id: Uuid, name: String, gender: String, artist_type: PersonType, tags: Vec<String>, release_groups: Vec<ReleaseGroup>) -> Artist {
        Artist {
            id: id,
            name: name,
            gender: gender,
            artist_type: artist_type,
            tags: tags,
            release_groups: release_groups
        }
    }
}

/// Provides methods for browsing, looking up or searching artists.
pub trait ArtistTrait {
    fn search(&self, params: &mut HashMap<&str, &str>) -> Vec<Artist>;
    fn lookup(&self, artist: Artist, params: &mut HashMap<&str, &str>) -> Result<Artist, String>;
}

impl ArtistTrait for super::MusicBrainz {
    /// Searches MusicBrainz for artists based on the search query.
    ///
    /// Returns a `Vec` containing the artists matching the search query.
    /// If no artists were found, returns an empty `Vec`.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use musicbrainz::*;
    /// let musicbrainz = MusicBrainz::new();
    /// let search_results = musicbrainz.search("deadmau5");
    ///
    /// assert_eq!(search_results[0].id.hyphenated().to_string(), "4a00ec9d-c635-463a-8cd4-eb61725f0c60");
    /// ```
    fn search(&self, params: &mut HashMap<&str, &str>) -> Vec<Artist> {
        params.insert("fmt", "json");
        let data = self.get("artist", params).unwrap();

        let count = data["count"].as_i32().unwrap();

        if count == 0 {
            return Vec::new();
        }

        let artists = &data["artists"];
        let mut results: Vec<Artist> = Vec::new();

        for artist in artists.members() {
            if !artist["score"].is_null() {
                if artist["score"] == "100" {
                    if !artist["name"].is_null() {
                        let name = artist["name"].to_string();
                        let gender = artist["gender"].to_string();
                        let id = Uuid::parse_str(artist["id"].as_str()
                            .expect("failed to parse artist ID as slice"))
                            .expect("failed to parse artist ID as Uuid");
                        let mut tags: Vec<String> = Vec::new();
                        let release_groups: Vec<ReleaseGroup> = Vec::new();

                        for tag in artist["tags"].members() {
                            tags.push(tag["name"].to_string());
                        }

                        results.push(Artist::new(id, name, gender, PersonType::Other, tags, release_groups));
                    }
                }
            }
        }
        results
    }

    /// Lookup an artist by using its MusicBrainz Identifier.
    ///
    fn lookup(&self, artist: Artist, params: &mut HashMap<&str, &str>) -> Result<Artist, String> {
        params.insert("fmt", "json");

        let artist_clone = artist.clone();
        let id = artist_clone.id.hyphenated().to_string();
        let artist_data = self.get(&format!("artist/{id}", id=&id), params).unwrap();

        if !artist_data["error"].is_null() {
            let error_msg = artist_data["error"].to_string();
            return Err(format!("error looking up artist: {msg}", msg=error_msg));
        }

        let artist_type = artist_data["type"].as_str().expect("failed to parse artist type as slice").parse::<PersonType>().unwrap();
        let mut artist_albums: Vec<ReleaseGroup> = Vec::new();

        let mut tags: Vec<String> = Vec::new();
        if !artist_data["tags"].is_null() {
            for tag in artist_data["tags"].members() {
                tags.push(tag["name"].to_string());
            }
        }

        let release_groups = &artist_data["release-groups"];
        for album in release_groups.members() {
            let mut secondary_types: Vec<AlbumSecondaryType> = Vec::new();
            for secondary_type in album["secondary-types"].members() {
                secondary_types.push(secondary_type.as_str()
                    .expect("failed to parse album secondary type as slice")
                    .parse::<AlbumSecondaryType>()
                    .unwrap())
            }

            artist_albums.push(ReleaseGroup::new(
                album["title"].to_string(),
                album["first-release-date"].to_string(),
                Uuid::parse_str(album["id"].as_str()
                    .expect("failed to parse release group ID as slice"))
                    .expect("failed to parse release group ID as Uuid"),
                Uuid::parse_str(&artist_data["id"].as_str()
                    .expect("failed to parse artist ID as slice"))
                    .expect("failed to parse artist ID as Uuid"),
                album["primary-type"].as_str()
                    .expect("failed to parse album primary type as slice")
                    .parse::<AlbumMainType>()
                    .unwrap(),
                secondary_types
            ));
        }

        Ok(Artist::new(
            Uuid::parse_str(artist_data["id"].as_str()
                    .expect("failed to parse artist ID as slice"))
                    .expect("failed to parse artist ID as Uuid"),
            artist_data["name"].to_string(),
            artist_data["gender"].to_string(),
            artist_type,
            tags,
            artist_albums
        ))
    }
}
<commit_msg>Add entity type to trait methods<commit_after>use release_group::ReleaseGroup;
use uuid::Uuid;
use enums::{PersonType, AlbumMainType, AlbumSecondaryType};
use std::collections::HashMap;

#[derive(Debug, Clone)]
pub struct Artist {
    pub id: Uuid,
    pub name: String,
    pub gender: String,
    pub artist_type: PersonType,
    pub tags: Vec<String>,
    pub release_groups: Vec<ReleaseGroup>
}

impl Artist {
    pub fn new(id: Uuid, name: String, gender: String, artist_type: PersonType, tags: Vec<String>, release_groups: Vec<ReleaseGroup>) -> Artist {
        Artist {
            id: id,
            name: name,
            gender: gender,
            artist_type: artist_type,
            tags: tags,
            release_groups: release_groups
        }
    }
}

/// Provides methods for browsing, looking up or searching artists.
pub trait ArtistTrait {
    fn search_artist(&self, params: &mut HashMap<&str, &str>) -> Vec<Artist>;
    fn lookup_artist(&self, artist: Artist, params: &mut HashMap<&str, &str>) -> Result<Artist, String>;
}

impl ArtistTrait for super::MusicBrainz {
    /// Searches MusicBrainz for artists based on the search query.
    ///
    /// Returns a `Vec` containing the artists matching the search query.
    /// If no artists were found, returns an empty `Vec`.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use musicbrainz::*;
    /// let musicbrainz = MusicBrainz::new();
    /// let search_results = musicbrainz.search_artist("deadmau5");
    ///
    /// assert_eq!(search_results[0].id.hyphenated().to_string(), "4a00ec9d-c635-463a-8cd4-eb61725f0c60");
    /// ```
    fn search_artist(&self, params: &mut HashMap<&str, &str>) -> Vec<Artist> {
        params.insert("fmt", "json");
        let data = self.get("artist", params).unwrap();

        let count = data["count"].as_i32().unwrap();

        if count == 0 {
            return Vec::new();
        }

        let artists = &data["artists"];
        let mut results: Vec<Artist> = Vec::new();

        for artist in artists.members() {
            if !artist["score"].is_null() {
                if artist["score"] == "100" {
                    if !artist["name"].is_null() {
                        let name = artist["name"].to_string();
                        let gender = artist["gender"].to_string();
                        let id = Uuid::parse_str(artist["id"].as_str()
                            .expect("failed to parse artist ID as slice"))
                            .expect("failed to parse artist ID as Uuid");
                        let mut tags: Vec<String> = Vec::new();
                        let release_groups: Vec<ReleaseGroup> = Vec::new();

                        for tag in artist["tags"].members() {
                            tags.push(tag["name"].to_string());
                        }

                        results.push(Artist::new(id, name, gender, PersonType::Other, tags, release_groups));
                    }
                }
            }
        }
        results
    }

    /// Lookup an artist by using its MusicBrainz Identifier.
    ///
    fn lookup_artist(&self, artist: Artist, params: &mut HashMap<&str, &str>) -> Result<Artist, String> {
        params.insert("fmt", "json");

        let artist_clone = artist.clone();
        let id = artist_clone.id.hyphenated().to_string();
        let artist_data = self.get(&format!("artist/{id}", id=&id), params).unwrap();

        if !artist_data["error"].is_null() {
            let error_msg = artist_data["error"].to_string();
            return Err(format!("error looking up artist: {msg}", msg=error_msg));
        }

        let artist_type = artist_data["type"].as_str().expect("failed to parse artist type as slice").parse::<PersonType>().unwrap();
        let mut artist_albums: Vec<ReleaseGroup> = Vec::new();

        let mut tags: Vec<String> = Vec::new();
        if !artist_data["tags"].is_null() {
            for tag in artist_data["tags"].members() {
                tags.push(tag["name"].to_string());
            }
        }

        let release_groups = &artist_data["release-groups"];
        for album in release_groups.members() {
            let mut secondary_types: Vec<AlbumSecondaryType> = Vec::new();
            for secondary_type in album["secondary-types"].members() {
                secondary_types.push(secondary_type.as_str()
                    .expect("failed to parse album secondary type as slice")
                    .parse::<AlbumSecondaryType>()
                    .unwrap())
            }

            artist_albums.push(ReleaseGroup::new(
                album["title"].to_string(),
                album["first-release-date"].to_string(),
                Uuid::parse_str(album["id"].as_str()
                    .expect("failed to parse release group ID as slice"))
                    .expect("failed to parse release group ID as Uuid"),
                Uuid::parse_str(&artist_data["id"].as_str()
                    .expect("failed to parse artist ID as slice"))
                    .expect("failed to parse artist ID as Uuid"),
                album["primary-type"].as_str()
                    .expect("failed to parse album primary type as slice")
                    .parse::<AlbumMainType>()
                    .unwrap(),
                secondary_types
            ));
        }

        Ok(Artist::new(
            Uuid::parse_str(artist_data["id"].as_str()
                    .expect("failed to parse artist ID as slice"))
                    .expect("failed to parse artist ID as Uuid"),
            artist_data["name"].to_string(),
            artist_data["gender"].to_string(),
            artist_type,
            tags,
            artist_albums
        ))
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::io::Read;

use grid::Grid;
use self::error::SourceError;

pub mod error;

impl<T: ?Sized + Read> Source for T {}

/// The `Source` trait allows to use any implementor of the `Read` trait
/// as an input source for the grid string format with no additional effort.
pub trait Source: Read {
    /// Creates a `Grid` from a readable source.
    /// Reads from the source until EOF, parses the data as a string,
    /// then checks the array for size and legality and converts it to a `Grid`
    ///
    /// # Failure
    ///
    /// Returns an error if either the read failed,
    /// a character other than `0`, `1`, `.` or `\n` was found,
    /// or the if the array is invalid (empty or non-square) or illegal.
    /// If the read was successful and no unexpected character was found,
    /// the faulty array is returned as well.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use std::error::Error;
    /// # use std::io;
    /// # use std::io::Write;
    /// # use takuzu::Source;
    /// let grid = match io::stdin().source() {
    ///     Ok(grid) => grid,
    ///     Err(err) => {
    ///         write!(io::stderr(), "error: {}\n", err.description()).unwrap();
    ///         return
    ///     },
    /// };
    /// ```
    fn source(&mut self) -> Result<Grid, SourceError> {
        let buffer = {
            let mut buffer = String::new();
            try!(self.read_to_string(&mut buffer));
            buffer
        };
        buffer.parse().map_err(|err| SourceError::Parsing(err))
    }
}
<commit_msg>enum constructor is FnOnce<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::io::Read;

use grid::Grid;
use self::error::SourceError;

pub mod error;

impl<T: ?Sized + Read> Source for T {}

/// The `Source` trait allows to use any implementor of the `Read` trait
/// as an input source for the grid string format with no additional effort.
pub trait Source: Read {
    /// Creates a `Grid` from a readable source.
    /// Reads from the source until EOF, parses the data as a string,
    /// then checks the array for size and legality and converts it to a `Grid`
    ///
    /// # Failure
    ///
    /// Returns an error if either the read failed,
    /// a character other than `0`, `1`, `.` or `\n` was found,
    /// or the if the array is invalid (empty or non-square) or illegal.
    /// If the read was successful and no unexpected character was found,
    /// the faulty array is returned as well.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use std::error::Error;
    /// # use std::io;
    /// # use std::io::Write;
    /// # use takuzu::Source;
    /// let grid = match io::stdin().source() {
    ///     Ok(grid) => grid,
    ///     Err(err) => {
    ///         write!(io::stderr(), "error: {}\n", err.description()).unwrap();
    ///         return
    ///     },
    /// };
    /// ```
    fn source(&mut self) -> Result<Grid, SourceError> {
        let buffer = {
            let mut buffer = String::new();
            try!(self.read_to_string(&mut buffer));
            buffer
        };
        buffer.parse().map_err(SourceError::Parsing)
    }
}
<|endoftext|>


<commit_before>use std::fmt::{self, Display, Debug};
use std::cell::Cell;
use std::io::Read;
use std::io;
use nom;

const DEFAULT_BUF_SIZE: usize = 8 * 1024;

pub trait StreamBuf<O> {
    fn fill(&mut self, min_count: usize) -> Result<(), io::Error>;
    fn recycle(&mut self);
    fn consume(&mut self, count: usize);
    fn data<'b>(&'b self) -> &'b[O];
    fn needed(&self) -> Option<nom::Needed>;
    fn apply<'b, C, CO, E>(&'b self, combinator: C) -> Result<Option<CO>, nom::Err<&'b [O], E>> where
        O: 'b, C: Fn(&'b [O]) -> nom::IResult<&'b [O], CO, E>;

    fn fill_apply<'b, C, CO, E>(&'b mut self, combinator: C) -> Result<Option<CO>, FillApplyError<&'b [O], E>> where
        C: Fn(&'b [O]) -> nom::IResult<&'b [O], CO, E> {

        match self.needed() {
            Some(nom::Needed::Size(bytes)) => try!(self.fill(bytes)),
            Some(nom::Needed::Unknown) => try!(self.fill(1)),
            None => ()
        }

        Ok(try!(self.apply(combinator)))
    }
}

#[derive(Debug)]
pub enum FillApplyError<I, E> {
    Io(io::Error),
    Parser(nom::Err<I, E>),
}

impl<I, E> From<io::Error> for FillApplyError<I, E> {
    fn from(e: io::Error) -> FillApplyError<I, E> {
        FillApplyError::Io(e)
    }
}

impl<I, E> From<nom::Err<I, E>> for FillApplyError<I, E> {
    fn from(e: nom::Err<I, E>) -> FillApplyError<I, E> {
        FillApplyError::Parser(e)
    }
}

impl<I, E> Display for FillApplyError<I, E> where I: Debug, E: Debug {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &FillApplyError::Io(ref e) => write!(f, "Failed to read data (IO error): {}", e),
            &FillApplyError::Parser(ref e) => write!(f, "Failed to parse data (parser error): {}", e),
        }
    }
}

pub struct ReadStreamBuf<R: Read> {
    reader: R,
    //TODO: alignment?
    buf: Vec<u8>,
    cap: usize,
    needed: Cell<Option<nom::Needed>>,
    offset: Cell<usize>,
}

impl<R: Read> ReadStreamBuf<R> {
    pub fn new(reader: R) -> ReadStreamBuf<R> {
        ReadStreamBuf::with_capacity(reader, DEFAULT_BUF_SIZE)
    }

    pub fn with_capacity(reader: R, cap: usize) -> ReadStreamBuf<R> {
        ReadStreamBuf {
            reader: reader,
            buf: Vec::with_capacity(cap),
            cap: cap,
            needed: Cell::new(Some(nom::Needed::Unknown)),
            offset: Cell::new(0),
        }
    }
}

impl<R: Read> StreamBuf<u8> for ReadStreamBuf<R> {
    fn fill(&mut self, min_bytes: usize) -> Result<(), io::Error> {
        let len = self.buf.len();
        let have = len - self.offset.get();

        if have >= min_bytes {
            return Ok(())
        }

        let needed = min_bytes - have;

        //TODO: enforce cap
        self.buf.resize(self.cap, 0);
        //trace!("fill needed: {}", needed);
        //trace!("buf write: {}..{} ({}); have: {} will have: {}", len, len + needed, self.buf[len..len + needed].len(), have, have + needed);
        if let Err(err) = self.reader.read_exact(&mut self.buf[len..len + needed]) {
            self.buf.resize(len, 0);
            return Err(err);
        }

        // Try to read to the end of the buffer if we can
        match self.reader.read(&mut self.buf[len + needed..self.cap]) {
            Err(err) => {
                self.buf.resize(len + needed, 0);
                return Err(err)
            },
            Ok(bytes_read) => {
                self.buf.resize(len + needed + bytes_read, 0);
            }
        }

        //trace!("buf have: {:?}", self.data());
        Ok(())
    }

    fn recycle(&mut self) {
        if self.offset.get() == 0 {
            return
        }
        self.buf = self.buf.split_off(self.offset.get());
        self.offset.set(0);
    }

    fn consume(&mut self, bytes: usize) {
        let len = self.buf.len();

        let consume = if self.offset.get() + bytes > len {
            len - self.offset.get()
        } else {
            bytes
        };
        self.offset.set(self.offset.get() + consume);
    }

    fn data<'b>(&'b self) -> &'b[u8] {
        &self.buf[self.offset.get()..self.buf.len()]
    }

    fn needed(&self) -> Option<nom::Needed> {
        self.needed.get().clone()
    }

    fn apply<'b, C, CO, E>(&'b self, combinator: C) -> Result<Option<CO>, nom::Err<&'b [u8], E>> where
        C: Fn(&'b [u8]) -> nom::IResult<&'b [u8], CO, E> {
        let data = self.data();
        match combinator(data) {
            nom::IResult::Done(left, out) => {
                let consumed = data.len() - left.len();
                //trace!("done: consumed: {}", consumed);

                // Move the offset
                self.offset.set(self.offset.get() + consumed);
                // We don't know how much we will need now
                self.needed.set(None);
                Ok(Some(out))
            },
            nom::IResult::Error(err) => Err(err),
            nom::IResult::Incomplete(needed) => {
                //trace!("incomplete: needed: {:?}", needed);
                self.needed.set(Some(needed));
                Ok(None)
            }
        }
    }
}

#[cfg(test)]
mod resd_stream_buf_tests {
    use super::StreamBuf;
    use super::ReadStreamBuf;
    use nom::{IResult, Needed};
    use std::io::Cursor;

    fn subject(data: Vec<u8>) -> ReadStreamBuf<Cursor<Vec<u8>>> {
        ReadStreamBuf::new(Cursor::new(data))
    }

    fn subject_with_default_data() -> ReadStreamBuf<Cursor<Vec<u8>>> {
        subject(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
    }

    #[test]
    fn reading() {
        let mut rsb = subject_with_default_data();
        assert_eq!(rsb.data(), [].as_ref());

        // fill reads as much as it can non-blocking
        // TODO: how to test blocking streaming?
        rsb.fill(3).unwrap();
        assert_eq!(rsb.data(), [0, 1, 2, 3, 4, 5, 6, 7, 8 ,9].as_ref());
    }

    #[test]
    fn recycle() {
        let mut rsb = subject_with_default_data();

        //TODO: this test does nothing...
        rsb.fill(10).unwrap();
        rsb.recycle();
        assert_eq!(rsb.data(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9].as_ref());
    }

    #[test]
    fn consume() {
        let mut rsb = subject_with_default_data();

        rsb.fill(10).unwrap();
        rsb.consume(8);
        assert_eq!(rsb.data(), [8, 9].as_ref());
    }

    #[test]
    fn consume_more_than_we_have() {
        let mut rsb = subject_with_default_data();

        rsb.consume(3);
        rsb.fill(10).unwrap();
        assert_eq!(rsb.data(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9].as_ref());
    }

    #[test]
    fn apply_function() {
        use nom::be_u8;
        let mut rsb = subject_with_default_data();

        rsb.fill(1).unwrap();
        assert_eq!(rsb.apply(be_u8).unwrap(), Some(0));
    }

    #[test]
    fn apply_should_consume() {
        use nom::be_u8;
        let mut rsb = subject_with_default_data();

        rsb.fill(1).unwrap();
        assert_eq!(rsb.apply(be_u8).unwrap(), Some(0));
        rsb.fill(1).unwrap();
        assert_eq!(rsb.apply(be_u8).unwrap(), Some(1));
    }

    #[test]
    fn apply_converted_macro() {
        let mut rsb = subject_with_default_data();

        rsb.fill(2).unwrap();
        assert_eq!(rsb.apply(closure!(tag!([0, 1]))).unwrap(), Some([0, 1].as_ref()));
    }

    #[test]
    fn apply_closure() {
        let mut rsb = subject_with_default_data();

        rsb.fill(2).unwrap();
        assert_eq!(rsb.apply(|i| tag!(i, [0, 1])).unwrap(), Some([0, 1].as_ref()));
    }

    #[test]
    fn apply_custom_fuction_with_refs() {
        let mut rsb = subject_with_default_data();
        fn comb(input: &[u8]) -> IResult<&[u8], &[u8]> {
            tag!(input, [0, 1, 2])
        }

        rsb.fill(3).unwrap();
        assert_eq!(rsb.apply(comb).unwrap(), Some([0, 1, 2].as_ref()));
    }

    #[test]
    fn needed_with_apply() {
        let mut rsb = subject_with_default_data();

        fn comb<'a>(input: &'a[u8]) -> IResult<&'a[u8], &'a[u8]> {
            tag!(input, [0, 1, 2])
        }

        assert!(rsb.apply(comb).unwrap().is_none());

        let needed = rsb.needed().unwrap();
        assert_eq!(needed, Needed::Size(3));

        if let Needed::Size(bytes) = needed {
            rsb.fill(bytes).unwrap();
            assert_eq!(rsb.apply(comb).unwrap(), Some([0, 1, 2].as_ref()));
        }
    }

    #[test]
    fn fill_apply() {
        use nom::be_u8;
        let mut rsb = subject_with_default_data();

        fn comb(input: &[u8]) -> IResult<&[u8], &[u8]> {
            tag!(input, [0, 1, 2])
        }

        //assert_eq!(rsb.fill_apply(comb).unwrap(), None);
        assert_eq!(rsb.fill_apply(comb).unwrap(), Some([0, 1, 2].as_ref()));
        //assert_eq!(rsb.fill_apply(be_u8).unwrap(), None);
        assert_eq!(rsb.fill_apply(be_u8).unwrap(), Some(3));
    }
}
<commit_msg>handling of buffer recycle and max capacity panic<commit_after>use std::fmt::{self, Display, Debug};
use std::cell::Cell;
use std::io::Read;
use std::io;
use nom;

const DEFAULT_BUF_SIZE: usize = 8 * 1024;

pub trait StreamBuf<O> {
    fn fill(&mut self, min_count: usize) -> Result<(), io::Error>;
    fn recycle(&mut self);
    fn consume(&mut self, count: usize);
    fn data<'b>(&'b self) -> &'b[O];
    fn needed(&self) -> Option<nom::Needed>;
    fn apply<'b, C, CO, E>(&'b self, combinator: C) -> Result<Option<CO>, nom::Err<&'b [O], E>> where
        O: 'b, C: Fn(&'b [O]) -> nom::IResult<&'b [O], CO, E>;

    fn fill_apply<'b, C, CO, E>(&'b mut self, combinator: C) -> Result<Option<CO>, FillApplyError<&'b [O], E>> where
        C: Fn(&'b [O]) -> nom::IResult<&'b [O], CO, E> {

        match self.needed() {
            Some(nom::Needed::Size(bytes)) => try!(self.fill(bytes)),
            Some(nom::Needed::Unknown) => try!(self.fill(1)),
            None => ()
        }

        Ok(try!(self.apply(combinator)))
    }
}

#[derive(Debug)]
pub enum FillApplyError<I, E> {
    Io(io::Error),
    Parser(nom::Err<I, E>),
}

impl<I, E> From<io::Error> for FillApplyError<I, E> {
    fn from(e: io::Error) -> FillApplyError<I, E> {
        FillApplyError::Io(e)
    }
}

impl<I, E> From<nom::Err<I, E>> for FillApplyError<I, E> {
    fn from(e: nom::Err<I, E>) -> FillApplyError<I, E> {
        FillApplyError::Parser(e)
    }
}

impl<I, E> Display for FillApplyError<I, E> where I: Debug, E: Debug {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &FillApplyError::Io(ref e) => write!(f, "Failed to read data (IO error): {}", e),
            &FillApplyError::Parser(ref e) => write!(f, "Failed to parse data (parser error): {}", e),
        }
    }
}

pub struct ReadStreamBuf<R: Read> {
    reader: R,
    //TODO: alignment?
    buf: Vec<u8>,
    cap: usize,
    needed: Cell<Option<nom::Needed>>,
    offset: Cell<usize>,
}

impl<R: Read> ReadStreamBuf<R> {
    pub fn new(reader: R) -> ReadStreamBuf<R> {
        ReadStreamBuf::with_capacity(reader, DEFAULT_BUF_SIZE)
    }

    pub fn with_capacity(reader: R, cap: usize) -> ReadStreamBuf<R> {
        ReadStreamBuf {
            reader: reader,
            buf: Vec::with_capacity(cap),
            cap: cap,
            needed: Cell::new(Some(nom::Needed::Unknown)),
            offset: Cell::new(0),
        }
    }
}

impl<R: Read> StreamBuf<u8> for ReadStreamBuf<R> {
    fn fill(&mut self, min_bytes: usize) -> Result<(), io::Error> {
        let len = self.buf.len();
        let have = len - self.offset.get();

        if have >= min_bytes {
            return Ok(())
        }

        let needed = min_bytes - have;

        if min_bytes > self.cap {
            panic!("Cannot fit {} B of data in {} B buffer", min_bytes, self.cap)
        }
        if len + needed > self.cap {
            self.recycle();
            assert_eq!(self.offset.get(), 0);
            assert_eq!(self.buf.len(), have);
            return self.fill(min_bytes)
        }

        self.buf.resize(self.cap, 0);
        //trace!("fill needed: {}", needed);
        //trace!("buf write: {}..{} ({}); have: {} will have: {}", len, len + needed, self.buf[len..len + needed].len(), have, have + needed);
        if let Err(err) = self.reader.read_exact(&mut self.buf[len..len + needed]) {
            self.buf.resize(len, 0);
            return Err(err);
        }

        // Try to read to the end of the buffer if we can
        match self.reader.read(&mut self.buf[len + needed..self.cap]) {
            Err(err) => {
                self.buf.resize(len + needed, 0);
                return Err(err)
            },
            Ok(bytes_read) => {
                self.buf.resize(len + needed + bytes_read, 0);
            }
        }

        //trace!("buf have: {:?}", self.data());
        Ok(())
    }

    fn recycle(&mut self) {
        if self.offset.get() == 0 {
            return
        }
        self.buf = self.buf.split_off(self.offset.get());
        self.offset.set(0);
    }

    fn consume(&mut self, bytes: usize) {
        let len = self.buf.len();

        let consume = if self.offset.get() + bytes > len {
            len - self.offset.get()
        } else {
            bytes
        };
        self.offset.set(self.offset.get() + consume);
    }

    fn data<'b>(&'b self) -> &'b[u8] {
        &self.buf[self.offset.get()..self.buf.len()]
    }

    fn needed(&self) -> Option<nom::Needed> {
        self.needed.get().clone()
    }

    fn apply<'b, C, CO, E>(&'b self, combinator: C) -> Result<Option<CO>, nom::Err<&'b [u8], E>> where
        C: Fn(&'b [u8]) -> nom::IResult<&'b [u8], CO, E> {
        let data = self.data();
        match combinator(data) {
            nom::IResult::Done(left, out) => {
                let consumed = data.len() - left.len();
                //trace!("done: consumed: {}", consumed);

                // Move the offset
                self.offset.set(self.offset.get() + consumed);
                // We don't know how much we will need now
                self.needed.set(None);
                Ok(Some(out))
            },
            nom::IResult::Error(err) => Err(err),
            nom::IResult::Incomplete(needed) => {
                //trace!("incomplete: needed: {:?}", needed);
                self.needed.set(Some(needed));
                Ok(None)
            }
        }
    }
}

#[cfg(test)]
mod resd_stream_buf_tests {
    use super::StreamBuf;
    use super::ReadStreamBuf;
    use nom::{IResult, Needed};
    use std::io::Cursor;

    fn subject(data: Vec<u8>) -> ReadStreamBuf<Cursor<Vec<u8>>> {
        ReadStreamBuf::new(Cursor::new(data))
    }

    fn subject_with_default_data() -> ReadStreamBuf<Cursor<Vec<u8>>> {
        subject(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
    }

    #[test]
    fn reading() {
        let mut rsb = subject_with_default_data();
        assert_eq!(rsb.data(), [].as_ref());

        // fill reads as much as it can non-blocking
        // TODO: how to test blocking streaming?
        rsb.fill(3).unwrap();
        assert_eq!(rsb.data(), [0, 1, 2, 3, 4, 5, 6, 7, 8 ,9].as_ref());
    }

    #[test]
    fn recycle() {
        let mut rsb = subject_with_default_data();

        //TODO: this test does nothing...
        rsb.fill(10).unwrap();
        rsb.recycle();
        assert_eq!(rsb.data(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9].as_ref());
    }

    #[test]
    fn consume() {
        let mut rsb = subject_with_default_data();

        rsb.fill(10).unwrap();
        rsb.consume(8);
        assert_eq!(rsb.data(), [8, 9].as_ref());
    }

    #[test]
    fn consume_more_than_we_have() {
        let mut rsb = subject_with_default_data();

        rsb.consume(3);
        rsb.fill(10).unwrap();
        assert_eq!(rsb.data(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9].as_ref());
    }

    #[test]
    fn apply_function() {
        use nom::be_u8;
        let mut rsb = subject_with_default_data();

        rsb.fill(1).unwrap();
        assert_eq!(rsb.apply(be_u8).unwrap(), Some(0));
    }

    #[test]
    fn apply_should_consume() {
        use nom::be_u8;
        let mut rsb = subject_with_default_data();

        rsb.fill(1).unwrap();
        assert_eq!(rsb.apply(be_u8).unwrap(), Some(0));
        rsb.fill(1).unwrap();
        assert_eq!(rsb.apply(be_u8).unwrap(), Some(1));
    }

    #[test]
    fn apply_converted_macro() {
        let mut rsb = subject_with_default_data();

        rsb.fill(2).unwrap();
        assert_eq!(rsb.apply(closure!(tag!([0, 1]))).unwrap(), Some([0, 1].as_ref()));
    }

    #[test]
    fn apply_closure() {
        let mut rsb = subject_with_default_data();

        rsb.fill(2).unwrap();
        assert_eq!(rsb.apply(|i| tag!(i, [0, 1])).unwrap(), Some([0, 1].as_ref()));
    }

    #[test]
    fn apply_custom_fuction_with_refs() {
        let mut rsb = subject_with_default_data();
        fn comb(input: &[u8]) -> IResult<&[u8], &[u8]> {
            tag!(input, [0, 1, 2])
        }

        rsb.fill(3).unwrap();
        assert_eq!(rsb.apply(comb).unwrap(), Some([0, 1, 2].as_ref()));
    }

    #[test]
    fn needed_with_apply() {
        let mut rsb = subject_with_default_data();

        fn comb<'a>(input: &'a[u8]) -> IResult<&'a[u8], &'a[u8]> {
            tag!(input, [0, 1, 2])
        }

        assert!(rsb.apply(comb).unwrap().is_none());

        let needed = rsb.needed().unwrap();
        assert_eq!(needed, Needed::Size(3));

        if let Needed::Size(bytes) = needed {
            rsb.fill(bytes).unwrap();
            assert_eq!(rsb.apply(comb).unwrap(), Some([0, 1, 2].as_ref()));
        }
    }

    #[test]
    fn fill_apply() {
        use nom::be_u8;
        let mut rsb = subject_with_default_data();

        fn comb(input: &[u8]) -> IResult<&[u8], &[u8]> {
            tag!(input, [0, 1, 2])
        }

        //assert_eq!(rsb.fill_apply(comb).unwrap(), None);
        assert_eq!(rsb.fill_apply(comb).unwrap(), Some([0, 1, 2].as_ref()));
        //assert_eq!(rsb.fill_apply(be_u8).unwrap(), None);
        assert_eq!(rsb.fill_apply(be_u8).unwrap(), Some(3));
    }

    #[test]
    #[should_panic(expected = "Cannot fit 8193 B of data in 8192 B buffer")]
    fn fill_over_buf() {
        let mut rsb = subject_with_default_data();
        rsb.fill(super::DEFAULT_BUF_SIZE + 1).ok();
    }
}
<|endoftext|>


<commit_before>use crate::*;
use std::env;
use std::path;

#[test]
pub fn context_build_tests() {
    // TODO: More tests!
    let confs = vec! [
        conf::Conf::default()
            .window_mode(conf::WindowMode::default()
            .dimensions(800.0, 600.0)),
        conf::Conf::default()
            .window_mode(conf::WindowMode::default()
            .dimensions(400.0, 400.0)),
        conf::Conf::default()
            .window_mode(conf::WindowMode::default()
            .resizable(false)),
    ];
    for conf in confs.into_iter() {
        let mut cb = ContextBuilder::new("ggez_unit_tests", "ggez").conf(conf);
        if let Ok(manifest_dir) = env::var("CARGO_MANIFEST_DIR") {
            let mut path = path::PathBuf::from(manifest_dir);
            path.push("resources");
            cb = cb.add_resource_path(path);
        }
        let (c, _e) = cb.clone().build().unwrap();
        let (w, h) = graphics::drawable_size(&c);
        assert_eq!(w, cb.conf.window_mode.width.into());
        assert_eq!(h, cb.conf.window_mode.height.into());
        // Can't really test whether or not the window is resizable?
    }
}
<commit_msg>Mo' unit tests, mo' problems<commit_after>use crate::*;
use std::env;
use std::path;

#[test]
pub fn context_build_tests() {
    // TODO: More tests!
    let confs = vec! [
        conf::Conf::default()
            .window_mode(conf::WindowMode::default()
            .dimensions(800.0, 600.0)),
        conf::Conf::default()
            .window_mode(conf::WindowMode::default()
            .dimensions(400.0, 400.0)),
        conf::Conf::default()
            .window_mode(conf::WindowMode::default()
            .resizable(false)),
        // conf::Conf::default()
        //     .window_mode(conf::WindowMode::default()
        //     .fullscreen_type(conf::FullscreenType::Windowed)),
        // conf::Conf::default()
        //     .window_mode(conf::WindowMode::default()
        //     .fullscreen_type(conf::FullscreenType::True)),

        // TODO: This makes the window size wrong, it seems.  Why?
        // conf::Conf::default()
        //     .modules(conf::ModuleConf::default()
        //         .gamepad(false)
        //     ),
        conf::Conf::default()
            .modules(conf::ModuleConf::default()
                .audio(false)
            ),
    ];
    for conf in confs.into_iter() {
        let mut cb = ContextBuilder::new("ggez_unit_tests", "ggez").conf(conf);
        if let Ok(manifest_dir) = env::var("CARGO_MANIFEST_DIR") {
            let mut path = path::PathBuf::from(manifest_dir);
            path.push("resources");
            cb = cb.add_resource_path(path);
        }
        let (c, _e) = cb.clone().build().unwrap();
        let (w, h) = graphics::drawable_size(&c);
        assert_eq!(w, cb.conf.window_mode.width.into());
        assert_eq!(h, cb.conf.window_mode.height.into());
        // Can't really test whether or not the window is resizable?
    }
}
<|endoftext|>


<commit_before>use {AtomicDateTime, Run, Segment, Time, TimeSpan, TimeStamp, TimerPhase, TimingMethod};
use TimerPhase::*;
use comparison::personal_best;
use parking_lot::RwLock;
use std::sync::Arc;
use std::mem;

#[derive(Debug, Clone)]
pub struct Timer {
    run: Run,
    phase: TimerPhase,
    current_split_index: isize,
    current_timing_method: TimingMethod,
    current_comparison: String,
    attempt_started: Option<AtomicDateTime>,
    attempt_ended: Option<AtomicDateTime>,
    start_time: TimeStamp,
    start_time_with_offset: TimeStamp,
    // This gets adjusted after unpausing
    adjusted_start_time: TimeStamp,
    time_paused_at: TimeSpan,
    is_game_time_paused: bool,
    game_time_pause_time: Option<TimeSpan>,
    loading_times: Option<TimeSpan>,
}

pub type SharedTimer = Arc<RwLock<Timer>>;

quick_error! {
    #[derive(Debug)]
    pub enum CreationError {
        EmptyRun
    }
}

impl Timer {
    #[inline]
    pub fn new(mut run: Run) -> Result<Self, CreationError> {
        if run.is_empty() {
            return Err(CreationError::EmptyRun);
        }

        run.regenerate_comparisons();
        let now = TimeStamp::now();

        Ok(Timer {
            run: run,
            phase: NotRunning,
            current_split_index: -1,
            current_timing_method: TimingMethod::RealTime,
            current_comparison: personal_best::NAME.into(),
            attempt_started: None,
            attempt_ended: None,
            start_time: now,
            start_time_with_offset: now,
            adjusted_start_time: now,
            time_paused_at: TimeSpan::zero(),
            is_game_time_paused: false,
            game_time_pause_time: None,
            loading_times: None,
        })
    }

    pub fn into_shared(self) -> SharedTimer {
        Arc::new(RwLock::new(self))
    }

    pub fn replace_run(&mut self, run: Run, update_splits: bool) -> Result<Run, Run> {
        if run.is_empty() {
            return Err(run);
        }

        self.reset(update_splits);
        if !run.comparisons().any(|c| c == self.current_comparison) {
            self.current_comparison = personal_best::NAME.to_string();
        }

        Ok(mem::replace(&mut self.run, run))
    }

    pub fn set_run(&mut self, run: Run) -> Result<(), Run> {
        self.replace_run(run, false).map(|_| ())
    }

    #[inline]
    pub fn run(&self) -> &Run {
        &self.run
    }

    #[inline]
    pub fn current_phase(&self) -> TimerPhase {
        self.phase
    }

    pub fn current_time(&self) -> Time {
        let real_time = match self.phase {
            NotRunning => Some(self.run.offset()),
            Running => Some(TimeStamp::now() - self.adjusted_start_time),
            Paused => Some(self.time_paused_at),
            Ended => self.run.segments().last().unwrap().split_time().real_time,
        };

        let game_time = match self.phase {
            NotRunning => Some(self.run.offset()),
            Ended => self.run.segments().last().unwrap().split_time().game_time,
            _ => if self.is_game_time_paused() {
                self.game_time_pause_time
            } else {
                TimeSpan::option_sub(
                    real_time,
                    if self.is_game_time_initialized() {
                        Some(self.loading_times())
                    } else {
                        None
                    },
                )
            },
        };

        Time::new()
            .with_real_time(real_time)
            .with_game_time(game_time)
    }

    #[inline]
    pub fn current_timing_method(&self) -> TimingMethod {
        self.current_timing_method
    }

    #[inline]
    pub fn set_current_timing_method(&mut self, method: TimingMethod) {
        self.current_timing_method = method;
    }

    #[inline]
    pub fn current_comparison(&self) -> &str {
        &self.current_comparison
    }

    pub fn current_split(&self) -> Option<&Segment> {
        if self.current_split_index >= 0 {
            self.run.segments().get(self.current_split_index as usize)
        } else {
            None
        }
    }

    fn current_split_mut(&mut self) -> Option<&mut Segment> {
        if self.current_split_index >= 0 {
            self.run
                .segments_mut()
                .get_mut(self.current_split_index as usize)
        } else {
            None
        }
    }

    #[inline]
    pub fn current_split_index(&self) -> isize {
        self.current_split_index
    }

    pub fn start(&mut self) {
        if self.phase == NotRunning {
            self.phase = Running;
            self.current_split_index = 0;
            self.attempt_started = Some(AtomicDateTime::now());
            self.start_time = TimeStamp::now();
            self.start_time_with_offset = self.start_time - self.run.offset();
            self.adjusted_start_time = self.start_time_with_offset;
            self.time_paused_at = self.run.offset();
            self.uninitialize_game_time();
            self.run.start_next_run();

            // TODO OnStart
        }
    }

    pub fn split(&mut self) {
        let current_time = self.current_time();
        if self.phase == Running &&
            current_time
                .real_time
                .map_or(false, |t| t >= TimeSpan::zero())
        {
            self.current_split_mut()
                .unwrap()
                .set_split_time(current_time);
            self.current_split_index += 1;
            if self.run.len() as isize == self.current_split_index {
                self.phase = Ended;
                self.attempt_ended = Some(AtomicDateTime::now());
            }
            self.run.mark_as_changed();

            // TODO OnSplit
        }
    }

    pub fn split_or_start(&mut self) {
        if self.phase == NotRunning {
            self.start();
        } else {
            self.split();
        }
    }

    pub fn skip_split(&mut self) {
        if (self.phase == Running || self.phase == Paused) &&
            self.current_split_index < self.run.len() as isize - 1
        {
            self.current_split_mut().unwrap().clear_split_time();
            self.current_split_index += 1;
            self.run.mark_as_changed();

            // TODO OnSkipSplit
        }
    }

    pub fn undo_split(&mut self) {
        if self.phase != NotRunning && self.current_split_index > 0 {
            if self.phase == Ended {
                self.phase = Running;
            }
            self.current_split_index -= 1;
            self.current_split_mut().unwrap().clear_split_time();
            self.run.mark_as_changed();

            // TODO OnUndoSplit
        }
    }

    pub fn reset(&mut self, update_splits: bool) {
        if self.phase != NotRunning {
            if self.phase != Ended {
                self.attempt_ended = Some(AtomicDateTime::now());
            }
            self.unpause_game_time();
            self.set_loading_times(TimeSpan::zero());

            if update_splits {
                self.update_attempt_history();
                self.update_best_segments();
                self.update_pb_splits();
                self.update_segment_history();
            }

            self.reset_splits();
            self.run.fix_splits();
            self.run.regenerate_comparisons();
        }
    }

    fn reset_splits(&mut self) {
        self.phase = NotRunning;
        self.current_split_index = -1;

        // Reset Splits
        for segment in self.run.segments_mut() {
            segment.clear_split_time();
        }

        // TODO OnReset
    }

    pub fn pause(&mut self) {
        if self.phase == Running {
            self.time_paused_at = self.current_time().real_time.unwrap();
            self.phase = Paused;

            // TODO OnPause
        }
    }

    pub fn resume(&mut self) {
        if self.phase == Paused {
            self.adjusted_start_time = TimeStamp::now() - self.time_paused_at;
            self.phase = Running;

            // TODO OnResume
        }
    }

    pub fn toggle_pause(&mut self) {
        match self.phase {
            Running => self.pause(),
            Paused => self.resume(),
            _ => {}
        }
    }

    pub fn toggle_pause_or_start(&mut self) {
        match self.phase {
            Running => self.pause(),
            Paused => self.resume(),
            NotRunning => self.start(),
            _ => {}
        }
    }

    pub fn undo_all_pauses(&mut self) {
        match self.current_phase() {
            Paused => self.pause(),
            Ended => {
                let pause_time = Some(self.get_pause_time().unwrap_or_default());

                let split_time = self.run
                    .segments_mut()
                    .iter_mut()
                    .last()
                    .unwrap()
                    .split_time_mut();

                *split_time += Time::new()
                    .with_real_time(pause_time)
                    .with_game_time(pause_time);
            }
            _ => {}
        }

        self.adjusted_start_time = self.start_time_with_offset;

        // TODO OnUndoAllPauses
    }

    pub fn switch_to_next_comparison(&mut self) {
        let mut comparisons = self.run.comparisons();
        let len = comparisons.len();
        let index = comparisons
            .position(|c| c == self.current_comparison)
            .unwrap();
        let index = (index + 1) % len;
        self.current_comparison = self.run.comparisons().nth(index).unwrap().to_owned();

        // TODO OnNextComparison
    }

    pub fn switch_to_previous_comparison(&mut self) {
        let mut comparisons = self.run.comparisons();
        let len = comparisons.len();
        let index = comparisons
            .position(|c| c == self.current_comparison)
            .unwrap();
        let index = (index + len - 1) % len;
        self.current_comparison = self.run.comparisons().nth(index).unwrap().to_owned();

        // TODO OnPreviousComparison
    }

    pub fn current_attempt_duration(&self) -> TimeSpan {
        match self.current_phase() {
            NotRunning => TimeSpan::zero(),
            Paused | Running => TimeStamp::now() - self.start_time,
            Ended => self.attempt_ended.unwrap() - self.attempt_started.unwrap(),
        }
    }

    pub fn get_pause_time(&self) -> Option<TimeSpan> {
        match self.current_phase() {
            Paused => Some(
                TimeStamp::now() - self.start_time_with_offset - self.time_paused_at,
            ),
            Running | Ended if self.start_time_with_offset != self.adjusted_start_time => {
                Some(self.adjusted_start_time - self.start_time_with_offset)
            }
            _ => None,
        }
    }

    #[inline]
    pub fn is_game_time_initialized(&self) -> bool {
        self.loading_times.is_some()
    }

    #[inline]
    pub fn initialize_game_time(&mut self) {
        self.loading_times = Some(self.loading_times());
    }

    #[inline]
    pub fn uninitialize_game_time(&mut self) {
        self.loading_times = None;
    }

    #[inline]
    pub fn is_game_time_paused(&self) -> bool {
        self.is_game_time_paused
    }

    pub fn pause_game_time(&mut self) {
        if !self.is_game_time_paused() {
            let current_time = self.current_time();
            self.game_time_pause_time = current_time.game_time.or(current_time.real_time);
            self.is_game_time_paused = true;
        }
    }

    pub fn unpause_game_time(&mut self) {
        if self.is_game_time_paused() {
            let current_time = self.current_time();
            self.set_loading_times(
                TimeSpan::option_sub(current_time.real_time, current_time.game_time)
                    .unwrap_or_default(),
            );
            self.is_game_time_paused = false;
        }
    }

    #[inline]
    pub fn set_game_time(&mut self, game_time: TimeSpan) {
        if self.is_game_time_paused() {
            self.game_time_pause_time = Some(game_time);
        }
        let loading_times = self.current_time().real_time.unwrap() - game_time;
        self.loading_times = Some(loading_times);
    }

    #[inline]
    pub fn loading_times(&self) -> TimeSpan {
        self.loading_times.unwrap_or_default()
    }

    #[inline]
    pub fn set_loading_times(&mut self, time: TimeSpan) {
        self.loading_times = Some(time);
        if self.is_game_time_paused() {
            self.game_time_pause_time = Some(self.current_time().real_time.unwrap() - time);
        }
    }

    fn update_attempt_history(&mut self) {
        let time = if self.phase == Ended {
            self.current_time()
        } else {
            Default::default()
        };

        let pause_time = self.get_pause_time();

        self.run
            .add_attempt(time, self.attempt_started, self.attempt_ended, pause_time);
    }

    fn update_best_segments(&mut self) {
        let mut previous_split_time_rta = Some(TimeSpan::zero());
        let mut previous_split_time_game_time = Some(TimeSpan::zero());

        for split in self.run.segments_mut() {
            let mut new_best_segment = split.best_segment_time();
            if let Some(split_time) = split.split_time().real_time {
                let current_segment = previous_split_time_rta.map(|previous| split_time - previous);
                previous_split_time_rta = Some(split_time);
                if split
                    .best_segment_time()
                    .real_time
                    .map_or(true, |b| current_segment.map_or(false, |c| c < b))
                {
                    new_best_segment.real_time = current_segment;
                }
            }
            if let Some(split_time) = split.split_time().game_time {
                let current_segment =
                    previous_split_time_game_time.map(|previous| split_time - previous);
                previous_split_time_game_time = Some(split_time);
                if split
                    .best_segment_time()
                    .game_time
                    .map_or(true, |b| current_segment.map_or(false, |c| c < b))
                {
                    new_best_segment.game_time = current_segment;
                }
            }
            split.set_best_segment_time(new_best_segment);
        }
    }

    fn update_pb_splits(&mut self) {
        let method = self.current_timing_method;
        let (split_time, pb_split_time) = {
            let last_segment = self.run.segments().last().unwrap();
            (
                last_segment.split_time()[method],
                last_segment.personal_best_split_time()[method],
            )
        };
        if split_time.map_or(false, |s| pb_split_time.map_or(true, |pb| s < pb)) {
            self.set_run_as_pb();
        }
    }

    fn update_segment_history(&mut self) {
        self.run.update_segment_history(self.current_split_index);
    }

    fn set_run_as_pb(&mut self) {
        self.run.import_segment_history();
        self.run.fix_splits();
        for segment in self.run.segments_mut() {
            let split_time = segment.split_time();
            segment.set_personal_best_split_time(split_time);
        }
    }
}
<commit_msg>Fix Bug in Undo All Pauses<commit_after>use {AtomicDateTime, Run, Segment, Time, TimeSpan, TimeStamp, TimerPhase, TimingMethod};
use TimerPhase::*;
use comparison::personal_best;
use parking_lot::RwLock;
use std::sync::Arc;
use std::mem;

#[derive(Debug, Clone)]
pub struct Timer {
    run: Run,
    phase: TimerPhase,
    current_split_index: isize,
    current_timing_method: TimingMethod,
    current_comparison: String,
    attempt_started: Option<AtomicDateTime>,
    attempt_ended: Option<AtomicDateTime>,
    start_time: TimeStamp,
    start_time_with_offset: TimeStamp,
    // This gets adjusted after unpausing
    adjusted_start_time: TimeStamp,
    time_paused_at: TimeSpan,
    is_game_time_paused: bool,
    game_time_pause_time: Option<TimeSpan>,
    loading_times: Option<TimeSpan>,
}

pub type SharedTimer = Arc<RwLock<Timer>>;

quick_error! {
    #[derive(Debug)]
    pub enum CreationError {
        EmptyRun
    }
}

impl Timer {
    #[inline]
    pub fn new(mut run: Run) -> Result<Self, CreationError> {
        if run.is_empty() {
            return Err(CreationError::EmptyRun);
        }

        run.regenerate_comparisons();
        let now = TimeStamp::now();

        Ok(Timer {
            run: run,
            phase: NotRunning,
            current_split_index: -1,
            current_timing_method: TimingMethod::RealTime,
            current_comparison: personal_best::NAME.into(),
            attempt_started: None,
            attempt_ended: None,
            start_time: now,
            start_time_with_offset: now,
            adjusted_start_time: now,
            time_paused_at: TimeSpan::zero(),
            is_game_time_paused: false,
            game_time_pause_time: None,
            loading_times: None,
        })
    }

    pub fn into_shared(self) -> SharedTimer {
        Arc::new(RwLock::new(self))
    }

    pub fn replace_run(&mut self, run: Run, update_splits: bool) -> Result<Run, Run> {
        if run.is_empty() {
            return Err(run);
        }

        self.reset(update_splits);
        if !run.comparisons().any(|c| c == self.current_comparison) {
            self.current_comparison = personal_best::NAME.to_string();
        }

        Ok(mem::replace(&mut self.run, run))
    }

    pub fn set_run(&mut self, run: Run) -> Result<(), Run> {
        self.replace_run(run, false).map(|_| ())
    }

    #[inline]
    pub fn run(&self) -> &Run {
        &self.run
    }

    #[inline]
    pub fn current_phase(&self) -> TimerPhase {
        self.phase
    }

    pub fn current_time(&self) -> Time {
        let real_time = match self.phase {
            NotRunning => Some(self.run.offset()),
            Running => Some(TimeStamp::now() - self.adjusted_start_time),
            Paused => Some(self.time_paused_at),
            Ended => self.run.segments().last().unwrap().split_time().real_time,
        };

        let game_time = match self.phase {
            NotRunning => Some(self.run.offset()),
            Ended => self.run.segments().last().unwrap().split_time().game_time,
            _ => if self.is_game_time_paused() {
                self.game_time_pause_time
            } else {
                TimeSpan::option_sub(
                    real_time,
                    if self.is_game_time_initialized() {
                        Some(self.loading_times())
                    } else {
                        None
                    },
                )
            },
        };

        Time::new()
            .with_real_time(real_time)
            .with_game_time(game_time)
    }

    #[inline]
    pub fn current_timing_method(&self) -> TimingMethod {
        self.current_timing_method
    }

    #[inline]
    pub fn set_current_timing_method(&mut self, method: TimingMethod) {
        self.current_timing_method = method;
    }

    #[inline]
    pub fn current_comparison(&self) -> &str {
        &self.current_comparison
    }

    pub fn current_split(&self) -> Option<&Segment> {
        if self.current_split_index >= 0 {
            self.run.segments().get(self.current_split_index as usize)
        } else {
            None
        }
    }

    fn current_split_mut(&mut self) -> Option<&mut Segment> {
        if self.current_split_index >= 0 {
            self.run
                .segments_mut()
                .get_mut(self.current_split_index as usize)
        } else {
            None
        }
    }

    #[inline]
    pub fn current_split_index(&self) -> isize {
        self.current_split_index
    }

    pub fn start(&mut self) {
        if self.phase == NotRunning {
            self.phase = Running;
            self.current_split_index = 0;
            self.attempt_started = Some(AtomicDateTime::now());
            self.start_time = TimeStamp::now();
            self.start_time_with_offset = self.start_time - self.run.offset();
            self.adjusted_start_time = self.start_time_with_offset;
            self.time_paused_at = self.run.offset();
            self.uninitialize_game_time();
            self.run.start_next_run();

            // TODO OnStart
        }
    }

    pub fn split(&mut self) {
        let current_time = self.current_time();
        if self.phase == Running &&
            current_time
                .real_time
                .map_or(false, |t| t >= TimeSpan::zero())
        {
            self.current_split_mut()
                .unwrap()
                .set_split_time(current_time);
            self.current_split_index += 1;
            if self.run.len() as isize == self.current_split_index {
                self.phase = Ended;
                self.attempt_ended = Some(AtomicDateTime::now());
            }
            self.run.mark_as_changed();

            // TODO OnSplit
        }
    }

    pub fn split_or_start(&mut self) {
        if self.phase == NotRunning {
            self.start();
        } else {
            self.split();
        }
    }

    pub fn skip_split(&mut self) {
        if (self.phase == Running || self.phase == Paused) &&
            self.current_split_index < self.run.len() as isize - 1
        {
            self.current_split_mut().unwrap().clear_split_time();
            self.current_split_index += 1;
            self.run.mark_as_changed();

            // TODO OnSkipSplit
        }
    }

    pub fn undo_split(&mut self) {
        if self.phase != NotRunning && self.current_split_index > 0 {
            if self.phase == Ended {
                self.phase = Running;
            }
            self.current_split_index -= 1;
            self.current_split_mut().unwrap().clear_split_time();
            self.run.mark_as_changed();

            // TODO OnUndoSplit
        }
    }

    pub fn reset(&mut self, update_splits: bool) {
        if self.phase != NotRunning {
            if self.phase != Ended {
                self.attempt_ended = Some(AtomicDateTime::now());
            }
            self.unpause_game_time();
            self.set_loading_times(TimeSpan::zero());

            if update_splits {
                self.update_attempt_history();
                self.update_best_segments();
                self.update_pb_splits();
                self.update_segment_history();
            }

            self.reset_splits();
            self.run.fix_splits();
            self.run.regenerate_comparisons();
        }
    }

    fn reset_splits(&mut self) {
        self.phase = NotRunning;
        self.current_split_index = -1;

        // Reset Splits
        for segment in self.run.segments_mut() {
            segment.clear_split_time();
        }

        // TODO OnReset
    }

    pub fn pause(&mut self) {
        if self.phase == Running {
            self.time_paused_at = self.current_time().real_time.unwrap();
            self.phase = Paused;

            // TODO OnPause
        }
    }

    pub fn resume(&mut self) {
        if self.phase == Paused {
            self.adjusted_start_time = TimeStamp::now() - self.time_paused_at;
            self.phase = Running;

            // TODO OnResume
        }
    }

    pub fn toggle_pause(&mut self) {
        match self.phase {
            Running => self.pause(),
            Paused => self.resume(),
            _ => {}
        }
    }

    pub fn toggle_pause_or_start(&mut self) {
        match self.phase {
            Running => self.pause(),
            Paused => self.resume(),
            NotRunning => self.start(),
            _ => {}
        }
    }

    pub fn undo_all_pauses(&mut self) {
        match self.current_phase() {
            Paused => self.resume(),
            Ended => {
                let pause_time = Some(self.get_pause_time().unwrap_or_default());

                let split_time = self.run
                    .segments_mut()
                    .iter_mut()
                    .last()
                    .unwrap()
                    .split_time_mut();

                *split_time += Time::new()
                    .with_real_time(pause_time)
                    .with_game_time(pause_time);
            }
            _ => {}
        }

        self.adjusted_start_time = self.start_time_with_offset;

        // TODO OnUndoAllPauses
    }

    pub fn switch_to_next_comparison(&mut self) {
        let mut comparisons = self.run.comparisons();
        let len = comparisons.len();
        let index = comparisons
            .position(|c| c == self.current_comparison)
            .unwrap();
        let index = (index + 1) % len;
        self.current_comparison = self.run.comparisons().nth(index).unwrap().to_owned();

        // TODO OnNextComparison
    }

    pub fn switch_to_previous_comparison(&mut self) {
        let mut comparisons = self.run.comparisons();
        let len = comparisons.len();
        let index = comparisons
            .position(|c| c == self.current_comparison)
            .unwrap();
        let index = (index + len - 1) % len;
        self.current_comparison = self.run.comparisons().nth(index).unwrap().to_owned();

        // TODO OnPreviousComparison
    }

    pub fn current_attempt_duration(&self) -> TimeSpan {
        match self.current_phase() {
            NotRunning => TimeSpan::zero(),
            Paused | Running => TimeStamp::now() - self.start_time,
            Ended => self.attempt_ended.unwrap() - self.attempt_started.unwrap(),
        }
    }

    pub fn get_pause_time(&self) -> Option<TimeSpan> {
        match self.current_phase() {
            Paused => Some(
                TimeStamp::now() - self.start_time_with_offset - self.time_paused_at,
            ),
            Running | Ended if self.start_time_with_offset != self.adjusted_start_time => {
                Some(self.adjusted_start_time - self.start_time_with_offset)
            }
            _ => None,
        }
    }

    #[inline]
    pub fn is_game_time_initialized(&self) -> bool {
        self.loading_times.is_some()
    }

    #[inline]
    pub fn initialize_game_time(&mut self) {
        self.loading_times = Some(self.loading_times());
    }

    #[inline]
    pub fn uninitialize_game_time(&mut self) {
        self.loading_times = None;
    }

    #[inline]
    pub fn is_game_time_paused(&self) -> bool {
        self.is_game_time_paused
    }

    pub fn pause_game_time(&mut self) {
        if !self.is_game_time_paused() {
            let current_time = self.current_time();
            self.game_time_pause_time = current_time.game_time.or(current_time.real_time);
            self.is_game_time_paused = true;
        }
    }

    pub fn unpause_game_time(&mut self) {
        if self.is_game_time_paused() {
            let current_time = self.current_time();
            self.set_loading_times(
                TimeSpan::option_sub(current_time.real_time, current_time.game_time)
                    .unwrap_or_default(),
            );
            self.is_game_time_paused = false;
        }
    }

    #[inline]
    pub fn set_game_time(&mut self, game_time: TimeSpan) {
        if self.is_game_time_paused() {
            self.game_time_pause_time = Some(game_time);
        }
        let loading_times = self.current_time().real_time.unwrap() - game_time;
        self.loading_times = Some(loading_times);
    }

    #[inline]
    pub fn loading_times(&self) -> TimeSpan {
        self.loading_times.unwrap_or_default()
    }

    #[inline]
    pub fn set_loading_times(&mut self, time: TimeSpan) {
        self.loading_times = Some(time);
        if self.is_game_time_paused() {
            self.game_time_pause_time = Some(self.current_time().real_time.unwrap() - time);
        }
    }

    fn update_attempt_history(&mut self) {
        let time = if self.phase == Ended {
            self.current_time()
        } else {
            Default::default()
        };

        let pause_time = self.get_pause_time();

        self.run
            .add_attempt(time, self.attempt_started, self.attempt_ended, pause_time);
    }

    fn update_best_segments(&mut self) {
        let mut previous_split_time_rta = Some(TimeSpan::zero());
        let mut previous_split_time_game_time = Some(TimeSpan::zero());

        for split in self.run.segments_mut() {
            let mut new_best_segment = split.best_segment_time();
            if let Some(split_time) = split.split_time().real_time {
                let current_segment = previous_split_time_rta.map(|previous| split_time - previous);
                previous_split_time_rta = Some(split_time);
                if split
                    .best_segment_time()
                    .real_time
                    .map_or(true, |b| current_segment.map_or(false, |c| c < b))
                {
                    new_best_segment.real_time = current_segment;
                }
            }
            if let Some(split_time) = split.split_time().game_time {
                let current_segment =
                    previous_split_time_game_time.map(|previous| split_time - previous);
                previous_split_time_game_time = Some(split_time);
                if split
                    .best_segment_time()
                    .game_time
                    .map_or(true, |b| current_segment.map_or(false, |c| c < b))
                {
                    new_best_segment.game_time = current_segment;
                }
            }
            split.set_best_segment_time(new_best_segment);
        }
    }

    fn update_pb_splits(&mut self) {
        let method = self.current_timing_method;
        let (split_time, pb_split_time) = {
            let last_segment = self.run.segments().last().unwrap();
            (
                last_segment.split_time()[method],
                last_segment.personal_best_split_time()[method],
            )
        };
        if split_time.map_or(false, |s| pb_split_time.map_or(true, |pb| s < pb)) {
            self.set_run_as_pb();
        }
    }

    fn update_segment_history(&mut self) {
        self.run.update_segment_history(self.current_split_index);
    }

    fn set_run_as_pb(&mut self) {
        self.run.import_segment_history();
        self.run.fix_splits();
        for segment in self.run.segments_mut() {
            let split_time = segment.split_time();
            segment.set_personal_best_split_time(split_time);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the THIRD-PARTY file.

use std::{mem, result};

use super::get_fdt_addr;
use kvm_bindings::{
    user_pt_regs, KVM_REG_ARM64, KVM_REG_ARM64_SYSREG, KVM_REG_ARM64_SYSREG_CRM_MASK,
    KVM_REG_ARM64_SYSREG_CRM_SHIFT, KVM_REG_ARM64_SYSREG_CRN_MASK, KVM_REG_ARM64_SYSREG_CRN_SHIFT,
    KVM_REG_ARM64_SYSREG_OP0_MASK, KVM_REG_ARM64_SYSREG_OP0_SHIFT, KVM_REG_ARM64_SYSREG_OP1_MASK,
    KVM_REG_ARM64_SYSREG_OP1_SHIFT, KVM_REG_ARM64_SYSREG_OP2_MASK, KVM_REG_ARM64_SYSREG_OP2_SHIFT,
    KVM_REG_ARM_CORE, KVM_REG_SIZE_U64,
};
use kvm_ioctls::VcpuFd;

use vm_memory::GuestMemoryMmap;

/// Errors thrown while setting aarch64 registers.
#[derive(Debug)]
pub enum Error {
    /// Failed to set core register (PC, PSTATE or general purpose ones).
    SetCoreRegister(kvm_ioctls::Error),
    /// Failed to get a system register.
    GetSysRegister(kvm_ioctls::Error),
}
type Result<T> = result::Result<T, Error>;

#[allow(non_upper_case_globals)]
// PSR (Processor State Register) bits.
// Taken from arch/arm64/include/uapi/asm/ptrace.h.
const PSR_MODE_EL1h: u64 = 0x0000_0005;
const PSR_F_BIT: u64 = 0x0000_0040;
const PSR_I_BIT: u64 = 0x0000_0080;
const PSR_A_BIT: u64 = 0x0000_0100;
const PSR_D_BIT: u64 = 0x0000_0200;
// Taken from arch/arm64/kvm/inject_fault.c.
const PSTATE_FAULT_BITS_64: u64 = PSR_MODE_EL1h | PSR_A_BIT | PSR_F_BIT | PSR_I_BIT | PSR_D_BIT;

// Following are macros that help with getting the ID of a aarch64 core register.
// The core register are represented by the user_pt_regs structure. Look for it in
// arch/arm64/include/uapi/asm/ptrace.h.

// This macro gets the offset of a structure (i.e `str`) member (i.e `field`) without having
// an instance of that structure.
// It uses a null pointer to retrieve the offset to the field.
// Inspired by C solution: `#define offsetof(str, f) ((size_t)(&((str *)0)->f))`.
// Doing `offset__of!(user_pt_regs, pstate)` in our rust code will trigger the following:
// unsafe { &(*(0 as *const user_pt_regs)).pstate as *const _ as usize }
// The dereference expression produces an lvalue, but that lvalue is not actually read from,
// we're just doing pointer math on it, so in theory, it should safe.
macro_rules! offset__of {
    ($str:ty, $field:ident) => {
        unsafe { &(*(std::ptr::null::<$str>())).$field as *const _ as usize }
    };
}

macro_rules! arm64_core_reg {
    ($reg: tt) => {
        // As per `kvm_arm_copy_reg_indices`, the id of a core register can be obtained like this:
        // `const u64 core_reg = KVM_REG_ARM64 | KVM_REG_SIZE_U64 | KVM_REG_ARM_CORE | i`, where i is obtained with:
        // `for (i = 0; i < sizeof(struct kvm_regs) / sizeof(__u32); i++) {`
        // We are using here `user_pt_regs` since this structure contains the core register and it is at
        // the start of `kvm_regs`.
        // struct kvm_regs {
        //	struct user_pt_regs regs;	/* sp = sp_el0 */
        //
        //	__u64	sp_el1;
        //	__u64	elr_el1;
        //
        //	__u64	spsr[KVM_NR_SPSR];
        //
        //	struct user_fpsimd_state fp_regs;
        //};
        // struct user_pt_regs {
        //	__u64		regs[31];
        //	__u64		sp;
        //	__u64		pc;
        //	__u64		pstate;
        //};
        // In our implementation we need: pc, pstate and user_pt_regs->regs[0].
        KVM_REG_ARM64 as u64
            | KVM_REG_SIZE_U64 as u64
            | u64::from(KVM_REG_ARM_CORE)
            | ((offset__of!(user_pt_regs, $reg) / mem::size_of::<u32>()) as u64)
    };
}

// This macro computes the ID of a specific ARM64 system register similar to how
// the kernel C macro does.
// https://elixir.bootlin.com/linux/v4.20.17/source/arch/arm64/include/uapi/asm/kvm.h#L203
macro_rules! arm64_sys_reg {
    ($name: tt, $op0: tt, $op1: tt, $crn: tt, $crm: tt, $op2: tt) => {
        const $name: u64 = KVM_REG_ARM64 as u64
            | KVM_REG_SIZE_U64 as u64
            | KVM_REG_ARM64_SYSREG as u64
            | ((($op0 as u64) << KVM_REG_ARM64_SYSREG_OP0_SHIFT)
                & KVM_REG_ARM64_SYSREG_OP0_MASK as u64)
            | ((($op1 as u64) << KVM_REG_ARM64_SYSREG_OP1_SHIFT)
                & KVM_REG_ARM64_SYSREG_OP1_MASK as u64)
            | ((($crn as u64) << KVM_REG_ARM64_SYSREG_CRN_SHIFT)
                & KVM_REG_ARM64_SYSREG_CRN_MASK as u64)
            | ((($crm as u64) << KVM_REG_ARM64_SYSREG_CRM_SHIFT)
                & KVM_REG_ARM64_SYSREG_CRM_MASK as u64)
            | ((($op2 as u64) << KVM_REG_ARM64_SYSREG_OP2_SHIFT)
                & KVM_REG_ARM64_SYSREG_OP2_MASK as u64);
    };
}

// Constant imported from the Linux kernel:
// https://elixir.bootlin.com/linux/v4.20.17/source/arch/arm64/include/asm/sysreg.h#L135
arm64_sys_reg!(MPIDR_EL1, 3, 0, 0, 0, 5);

/// Configure core registers for a given CPU.
///
/// # Arguments
///
/// * `vcpu` - Structure for the VCPU that holds the VCPU's fd.
/// * `cpu_id` - Index of current vcpu.
/// * `boot_ip` - Starting instruction pointer.
/// * `mem` - Reserved DRAM for current VM.
pub fn setup_regs(vcpu: &VcpuFd, cpu_id: u8, boot_ip: u64, mem: &GuestMemoryMmap) -> Result<()> {
    // Get the register index of the PSTATE (Processor State) register.
    vcpu.set_one_reg(arm64_core_reg!(pstate), PSTATE_FAULT_BITS_64)
        .map_err(Error::SetCoreRegister)?;

    // Other vCPUs are powered off initially awaiting PSCI wakeup.
    if cpu_id == 0 {
        // Setting the PC (Processor Counter) to the current program address (kernel address).
        vcpu.set_one_reg(arm64_core_reg!(pc), boot_ip)
            .map_err(Error::SetCoreRegister)?;

        // Last mandatory thing to set -> the address pointing to the FDT (also called DTB).
        // "The device tree blob (dtb) must be placed on an 8-byte boundary and must
        // not exceed 2 megabytes in size." -> https://www.kernel.org/doc/Documentation/arm64/booting.txt.
        // We are choosing to place it the end of DRAM. See `get_fdt_addr`.
        vcpu.set_one_reg(arm64_core_reg!(regs), get_fdt_addr(mem) as u64)
            .map_err(Error::SetCoreRegister)?;
    }
    Ok(())
}

/// Read the MPIDR - Multiprocessor Affinity Register.
///
/// # Arguments
///
/// * `vcpu` - Structure for the VCPU that holds the VCPU's fd.
pub fn read_mpidr(vcpu: &VcpuFd) -> Result<u64> {
    vcpu.get_one_reg(MPIDR_EL1).map_err(Error::GetSysRegister)
}

#[cfg(test)]
mod tests {
    use super::*;
    use aarch64::{arch_memory_regions, layout};
    use kvm_ioctls::Kvm;

    #[test]
    fn test_setup_regs() {
        let kvm = Kvm::new().unwrap();
        let vm = kvm.create_vm().unwrap();
        let vcpu = vm.create_vcpu(0).unwrap();
        let regions = arch_memory_regions(layout::FDT_MAX_SIZE + 0x1000);
        let mem = GuestMemoryMmap::from_ranges(&regions).expect("Cannot initialize memory");

        match setup_regs(&vcpu, 0, 0x0, &mem).unwrap_err() {
            Error::SetCoreRegister(ref e) => assert_eq!(e.errno(), libc::ENOEXEC),
            _ => panic!("Expected to receive Error::SetCoreRegister"),
        }
        let mut kvi: kvm_bindings::kvm_vcpu_init = kvm_bindings::kvm_vcpu_init::default();
        vm.get_preferred_target(&mut kvi).unwrap();
        vcpu.vcpu_init(&kvi).unwrap();

        assert!(setup_regs(&vcpu, 0, 0x0, &mem).is_ok());
    }
    #[test]
    fn test_read_mpidr() {
        let kvm = Kvm::new().unwrap();
        let vm = kvm.create_vm().unwrap();
        let vcpu = vm.create_vcpu(0).unwrap();
        let mut kvi: kvm_bindings::kvm_vcpu_init = kvm_bindings::kvm_vcpu_init::default();
        vm.get_preferred_target(&mut kvi).unwrap();

        // Must fail when vcpu is not initialized yet.
        assert!(read_mpidr(&vcpu).is_err());

        vcpu.vcpu_init(&kvi).unwrap();
        assert_eq!(read_mpidr(&vcpu).unwrap(), 0x8000_0000);
    }
}
<commit_msg>aarch64: parametrizable macro for getting core reg id<commit_after>// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the THIRD-PARTY file.

use std::{mem, result};

use super::get_fdt_addr;
use kvm_bindings::{
    user_pt_regs, KVM_REG_ARM64, KVM_REG_ARM64_SYSREG, KVM_REG_ARM64_SYSREG_CRM_MASK,
    KVM_REG_ARM64_SYSREG_CRM_SHIFT, KVM_REG_ARM64_SYSREG_CRN_MASK, KVM_REG_ARM64_SYSREG_CRN_SHIFT,
    KVM_REG_ARM64_SYSREG_OP0_MASK, KVM_REG_ARM64_SYSREG_OP0_SHIFT, KVM_REG_ARM64_SYSREG_OP1_MASK,
    KVM_REG_ARM64_SYSREG_OP1_SHIFT, KVM_REG_ARM64_SYSREG_OP2_MASK, KVM_REG_ARM64_SYSREG_OP2_SHIFT,
    KVM_REG_ARM_CORE, KVM_REG_SIZE_U64,
};
use kvm_ioctls::VcpuFd;

use vm_memory::GuestMemoryMmap;

/// Errors thrown while setting aarch64 registers.
#[derive(Debug)]
pub enum Error {
    /// Failed to set core register (PC, PSTATE or general purpose ones).
    SetCoreRegister(kvm_ioctls::Error),
    /// Failed to get a system register.
    GetSysRegister(kvm_ioctls::Error),
}
type Result<T> = result::Result<T, Error>;

#[allow(non_upper_case_globals)]
// PSR (Processor State Register) bits.
// Taken from arch/arm64/include/uapi/asm/ptrace.h.
const PSR_MODE_EL1h: u64 = 0x0000_0005;
const PSR_F_BIT: u64 = 0x0000_0040;
const PSR_I_BIT: u64 = 0x0000_0080;
const PSR_A_BIT: u64 = 0x0000_0100;
const PSR_D_BIT: u64 = 0x0000_0200;
// Taken from arch/arm64/kvm/inject_fault.c.
const PSTATE_FAULT_BITS_64: u64 = PSR_MODE_EL1h | PSR_A_BIT | PSR_F_BIT | PSR_I_BIT | PSR_D_BIT;

// Following are macros that help with getting the ID of a aarch64 core register.
// The core register are represented by the user_pt_regs structure. Look for it in
// arch/arm64/include/uapi/asm/ptrace.h.

// This macro gets the offset of a structure (i.e `str`) member (i.e `field`) without having
// an instance of that structure.
// It uses a null pointer to retrieve the offset to the field.
// Inspired by C solution: `#define offsetof(str, f) ((size_t)(&((str *)0)->f))`.
// Doing `offset__of!(user_pt_regs, pstate)` in our rust code will trigger the following:
// unsafe { &(*(0 as *const user_pt_regs)).pstate as *const _ as usize }
// The dereference expression produces an lvalue, but that lvalue is not actually read from,
// we're just doing pointer math on it, so in theory, it should safe.
macro_rules! offset__of {
    ($str:ty, $field:ident) => {
        unsafe { &(*(std::ptr::null::<$str>())).$field as *const _ as usize }
    };
}

macro_rules! arm64_core_reg {
    ($reg: tt) => {
        // As per `kvm_arm_copy_reg_indices`, the id of a core register can be obtained like this:
        // `const u64 core_reg = KVM_REG_ARM64 | KVM_REG_SIZE_U64 | KVM_REG_ARM_CORE | i`, where i is obtained with:
        // `for (i = 0; i < sizeof(struct kvm_regs) / sizeof(__u32); i++) {`
        // We are using here `user_pt_regs` since this structure contains the core register and it is at
        // the start of `kvm_regs`.
        // struct kvm_regs {
        //	struct user_pt_regs regs;	/* sp = sp_el0 */
        //
        //	__u64	sp_el1;
        //	__u64	elr_el1;
        //
        //	__u64	spsr[KVM_NR_SPSR];
        //
        //	struct user_fpsimd_state fp_regs;
        //};
        // struct user_pt_regs {
        //	__u64		regs[31];
        //	__u64		sp;
        //	__u64		pc;
        //	__u64		pstate;
        //};
        // In our implementation we need: pc, pstate and user_pt_regs->regs[0].
        KVM_REG_ARM64 as u64
            | KVM_REG_SIZE_U64 as u64
            | u64::from(KVM_REG_ARM_CORE)
            | ((offset__of!(user_pt_regs, $reg) / mem::size_of::<u32>()) as u64)
    };
}

macro_rules! arm64_core_reg_id {
    ($size: tt, $offset: tt) => {
        // The core registers of an arm64 machine are represented
        // in kernel by the `kvm_regs` structure. This structure is a
        // mix of 32, 64 and 128 bit fields:
        // struct kvm_regs {
        //     struct user_pt_regs      regs;
        //
        //     __u64                    sp_el1;
        //     __u64                    elr_el1;
        //
        //     __u64                    spsr[KVM_NR_SPSR];
        //
        //     struct user_fpsimd_state fp_regs;
        // };
        // struct user_pt_regs {
        //     __u64 regs[31];
        //     __u64 sp;
        //     __u64 pc;
        //     __u64 pstate;
        // };
        // The id of a core register can be obtained like this:
        // offset = id & ~(KVM_REG_ARCH_MASK | KVM_REG_SIZE_MASK | KVM_REG_ARM_CORE). Thus,
        // id = KVM_REG_ARM64 | KVM_REG_SIZE_U64/KVM_REG_SIZE_U32/KVM_REG_SIZE_U128 | KVM_REG_ARM_CORE | offset
        // Obs! This macro intends to replace the deprecated limited one: arm64_core_reg but in an ulterior commit.
        KVM_REG_ARM64 as u64
            | u64::from(KVM_REG_ARM_CORE)
            | $size
            | (($offset / mem::size_of::<u32>()) as u64)
    };
}

// This macro computes the ID of a specific ARM64 system register similar to how
// the kernel C macro does.
// https://elixir.bootlin.com/linux/v4.20.17/source/arch/arm64/include/uapi/asm/kvm.h#L203
macro_rules! arm64_sys_reg {
    ($name: tt, $op0: tt, $op1: tt, $crn: tt, $crm: tt, $op2: tt) => {
        const $name: u64 = KVM_REG_ARM64 as u64
            | KVM_REG_SIZE_U64 as u64
            | KVM_REG_ARM64_SYSREG as u64
            | ((($op0 as u64) << KVM_REG_ARM64_SYSREG_OP0_SHIFT)
                & KVM_REG_ARM64_SYSREG_OP0_MASK as u64)
            | ((($op1 as u64) << KVM_REG_ARM64_SYSREG_OP1_SHIFT)
                & KVM_REG_ARM64_SYSREG_OP1_MASK as u64)
            | ((($crn as u64) << KVM_REG_ARM64_SYSREG_CRN_SHIFT)
                & KVM_REG_ARM64_SYSREG_CRN_MASK as u64)
            | ((($crm as u64) << KVM_REG_ARM64_SYSREG_CRM_SHIFT)
                & KVM_REG_ARM64_SYSREG_CRM_MASK as u64)
            | ((($op2 as u64) << KVM_REG_ARM64_SYSREG_OP2_SHIFT)
                & KVM_REG_ARM64_SYSREG_OP2_MASK as u64);
    };
}

// Constant imported from the Linux kernel:
// https://elixir.bootlin.com/linux/v4.20.17/source/arch/arm64/include/asm/sysreg.h#L135
arm64_sys_reg!(MPIDR_EL1, 3, 0, 0, 0, 5);

/// Configure core registers for a given CPU.
///
/// # Arguments
///
/// * `vcpu` - Structure for the VCPU that holds the VCPU's fd.
/// * `cpu_id` - Index of current vcpu.
/// * `boot_ip` - Starting instruction pointer.
/// * `mem` - Reserved DRAM for current VM.
pub fn setup_regs(vcpu: &VcpuFd, cpu_id: u8, boot_ip: u64, mem: &GuestMemoryMmap) -> Result<()> {
    // Get the register index of the PSTATE (Processor State) register.
    vcpu.set_one_reg(arm64_core_reg!(pstate), PSTATE_FAULT_BITS_64)
        .map_err(Error::SetCoreRegister)?;

    // Other vCPUs are powered off initially awaiting PSCI wakeup.
    if cpu_id == 0 {
        // Setting the PC (Processor Counter) to the current program address (kernel address).
        vcpu.set_one_reg(arm64_core_reg!(pc), boot_ip)
            .map_err(Error::SetCoreRegister)?;

        // Last mandatory thing to set -> the address pointing to the FDT (also called DTB).
        // "The device tree blob (dtb) must be placed on an 8-byte boundary and must
        // not exceed 2 megabytes in size." -> https://www.kernel.org/doc/Documentation/arm64/booting.txt.
        // We are choosing to place it the end of DRAM. See `get_fdt_addr`.
        vcpu.set_one_reg(arm64_core_reg!(regs), get_fdt_addr(mem) as u64)
            .map_err(Error::SetCoreRegister)?;
    }
    Ok(())
}

/// Read the MPIDR - Multiprocessor Affinity Register.
///
/// # Arguments
///
/// * `vcpu` - Structure for the VCPU that holds the VCPU's fd.
pub fn read_mpidr(vcpu: &VcpuFd) -> Result<u64> {
    vcpu.get_one_reg(MPIDR_EL1).map_err(Error::GetSysRegister)
}

#[cfg(test)]
mod tests {
    use super::*;
    use aarch64::{arch_memory_regions, layout};
    use kvm_ioctls::Kvm;

    #[test]
    fn test_setup_regs() {
        let kvm = Kvm::new().unwrap();
        let vm = kvm.create_vm().unwrap();
        let vcpu = vm.create_vcpu(0).unwrap();
        let regions = arch_memory_regions(layout::FDT_MAX_SIZE + 0x1000);
        let mem = GuestMemoryMmap::from_ranges(&regions).expect("Cannot initialize memory");

        match setup_regs(&vcpu, 0, 0x0, &mem).unwrap_err() {
            Error::SetCoreRegister(ref e) => assert_eq!(e.errno(), libc::ENOEXEC),
            _ => panic!("Expected to receive Error::SetCoreRegister"),
        }
        let mut kvi: kvm_bindings::kvm_vcpu_init = kvm_bindings::kvm_vcpu_init::default();
        vm.get_preferred_target(&mut kvi).unwrap();
        vcpu.vcpu_init(&kvi).unwrap();

        assert!(setup_regs(&vcpu, 0, 0x0, &mem).is_ok());
    }
    #[test]
    fn test_read_mpidr() {
        let kvm = Kvm::new().unwrap();
        let vm = kvm.create_vm().unwrap();
        let vcpu = vm.create_vcpu(0).unwrap();
        let mut kvi: kvm_bindings::kvm_vcpu_init = kvm_bindings::kvm_vcpu_init::default();
        vm.get_preferred_target(&mut kvi).unwrap();

        // Must fail when vcpu is not initialized yet.
        assert!(read_mpidr(&vcpu).is_err());

        vcpu.vcpu_init(&kvi).unwrap();
        assert_eq!(read_mpidr(&vcpu).unwrap(), 0x8000_0000);
    }
}
<|endoftext|>