<commit_before>//! Account commitment managment

use std::collections::hash_map::{self, HashMap};
use utils::address::Address;
use utils::bigint::{M256, U256};

use vm::Storage;
use vm::errors::{RequireError, CommitError};

#[derive(Debug, Clone)]
/// A single account commitment.
pub enum AccountCommitment {
    /// Full account commitment. The client that committed account
    /// should not change the account in other EVMs if it decides to
    /// accept the result.
    Full {
        nonce: M256,
        address: Address,
        balance: U256,
        code: Vec<u8>,
    },
    /// Commit only code of the account. The client can keep changing
    /// it in other EVMs if the code remains unchanged.
    Code {
        address: Address,
        code: Vec<u8>,
    },
    /// Commit a storage. Must be used given a full account.
    Storage {
        address: Address,
        index: M256,
        value: M256,
    },
}

impl AccountCommitment {
    /// Address of this account commitment.
    pub fn address(&self) -> Address {
        match self {
            &AccountCommitment::Full {
                address,
                ..
            } => address,
            &AccountCommitment::Code {
                address,
                ..
            } => address,
            &AccountCommitment::Storage {
                address,
                ..
            } => address,
        }
    }
}

#[derive(Debug, Clone)]
/// Represents an account. This is usually returned by the EVM.
pub enum Account {
    /// A full account. The client is expected to replace its own account state with this.
    Full {
        nonce: M256,
        address: Address,
        balance: U256,
        changing_storage: Storage,
        code: Vec<u8>,
    },
    /// Only balance is changed, and it is increasing for this address.
    IncreaseBalance(Address, U256),
    /// Only balance is changed, and it is decreasing for this address.
    DecreaseBalance(Address, U256),
    /// Remove an account, the client is expected to handle only
    /// removal without any ether transfer.
    Remove(Address),
    /// Create a new account.
    Create {
        nonce: M256,
        address: Address,
        balance: U256,
        storage: Storage,
        code: Vec<u8>,
    },
}

impl Account {
    /// Address of this account.
    pub fn address(&self) -> Address {
        match self {
            &Account::Full {
                address,
                ..
            } => address,
            &Account::IncreaseBalance(address, _) => address,
            &Account::DecreaseBalance(address, _) => address,
            &Account::Remove(address) => address,
            &Account::Create {
                address,
                ..
            } => address,
        }
    }
}

#[derive(Debug, Clone)]
/// A struct that manages the current account state for one EVM.
pub struct AccountState {
    accounts: HashMap<Address, Account>,
    codes: HashMap<Address, Vec<u8>>,
}

impl Default for AccountState {
    fn default() -> Self {
        Self {
            accounts: HashMap::new(),
            codes: HashMap::new(),
        }
    }
}

impl AccountState {
    /// Returns all accounts right now in this account state.
    pub fn accounts(&self) -> hash_map::Values<Address, Account> {
        self.accounts.values()
    }

    /// Returns Ok(()) if a full account is in this account
    /// state. Otherwise raise a `RequireError`.
    pub fn require(&self, address: Address) -> Result<(), RequireError> {
        match self.accounts.get(&address) {
            Some(&Account::Full { .. }) => return Ok(()),
            _ => return Err(RequireError::Account(address)),
        }
    }

    /// Returns Ok(()) if either a full account or a partial code
    /// account is in this account state. Otherwise raise a
    /// `RequireError`.
    pub fn require_code(&self, address: Address) -> Result<(), RequireError> {
        if self.codes.contains_key(&address) {
            return Ok(());
        }
        match self.accounts.get(&address) {
            Some(&Account::Full { .. }) => return Ok(()),
            _ => return Err(RequireError::AccountCode(address)),
        }
    }

    /// Returns Ok(()) if the storage exists in the VM. Otherwise
    /// raise a `RequireError`.
    pub fn require_storage(&self, address: Address, index: M256) -> Result<(), RequireError> {
        self.storage(address)?.read(index).and_then(|_| Ok(()))
    }

    /// Commit an account commitment into this account state.
    pub fn commit(&mut self, commitment: AccountCommitment) -> Result<(), CommitError> {
        match commitment {
            AccountCommitment::Full {
                nonce,
                address,
                balance,
                code
            } => {
                let account = if self.accounts.contains_key(&address) {
                    match self.accounts.remove(&address).unwrap() {
                        Account::Full { .. } => return Err(CommitError::AlreadyCommitted),
                        Account::Create { .. } => return Err(CommitError::AlreadyCommitted),
                        Account::Remove(address) => Account::Remove(address),
                        Account::IncreaseBalance(address, topup) => {
                            Account::Full {
                                nonce,
                                address,
                                balance: balance + topup,
                                changing_storage: Storage::new(address, true),
                                code,
                            }
                        },
                        Account::DecreaseBalance(address, withdraw) => {
                            Account::Full {
                                nonce,
                                address,
                                balance: balance - withdraw,
                                changing_storage: Storage::new(address, true),
                                code,
                            }
                        },
                    }
                } else {
                    Account::Full {
                        nonce,
                        address,
                        balance,
                        changing_storage: Storage::new(address, true),
                        code,
                    }
                };

                self.accounts.insert(address, account);
            },
            AccountCommitment::Code {
                address,
                code,
            } => {
                if self.accounts.contains_key(&address) || self.codes.contains_key(&address) {
                    return Err(CommitError::AlreadyCommitted);
                }

                self.codes.insert(address, code);
            },
            AccountCommitment::Storage {
                address,
                index,
                value
            } => {
                match self.accounts.get_mut(&address) {
                    Some(&mut Account::Full {
                        ref mut changing_storage,
                        ..
                    }) => {
                        changing_storage.commit(index, value)?;
                    },
                    _ => {
                        return Err(CommitError::InvalidCommitment);
                    },
                }
            }
        }
        Ok(())
    }

    /// Find code by its address in this account state. If the search
    /// failed, returns a `RequireError`.
    pub fn code(&self, address: Address) -> Result<&[u8], RequireError> {
        if self.codes.contains_key(&address) {
            return Ok(self.codes.get(&address).unwrap().as_slice());
        }

        if self.accounts.contains_key(&address) {
            match self.accounts.get(&address).unwrap() {
                &Account::Full {
                    ref code,
                    ..
                } => return Ok(code.as_slice()),
                &Account::Create {
                    ref code,
                    ..
                } => return Ok(code.as_slice()),
                &Account::Remove(_) => return Ok(&[]),
                &Account::IncreaseBalance(address, _) => return Err(RequireError::Account(address)),
                &Account::DecreaseBalance(address, _) => return Err(RequireError::Account(address)),
            }
        }

        return Err(RequireError::AccountCode(address));
    }

    /// Find nonce by its address in this account state. If the search
    /// failed, returns a `RequireError`.
    pub fn nonce(&self, address: Address) -> Result<M256, RequireError> {
        if self.accounts.contains_key(&address) {
            match self.accounts.get(&address).unwrap() {
                &Account::Full {
                    nonce,
                    ..
                } => return Ok(nonce),
                _ => (),
            }
        }

        return Err(RequireError::Account(address));
    }

    /// Find balance by its address in this account state. If the
    /// search failed, returns a `RequireError`.
    pub fn balance(&self, address: Address) -> Result<U256, RequireError> {
        if self.accounts.contains_key(&address) {
            match self.accounts.get(&address).unwrap() {
                &Account::Full {
                    balance,
                    ..
                } => return Ok(balance),
                _ => (),
            }
        }

        return Err(RequireError::Account(address));
    }

    /// Returns the storage of an account. If the account is not yet
    /// committed, returns a `RequireError`.
    pub fn storage(&self, address: Address) -> Result<&Storage, RequireError> {
        if self.accounts.contains_key(&address) {
            match self.accounts.get(&address).unwrap() {
                &Account::Full {
                    ref changing_storage,
                    ..
                } => return Ok(changing_storage),
                &Account::Create {
                    ref storage,
                    ..
                } => return Ok(storage),
                _ => (),
            }
        }

        return Err(RequireError::Account(address));
    }

    /// Returns the mutable storage of an account. If the account is
    /// not yet committed. returns a `RequireError`.
    pub fn storage_mut(&mut self, address: Address) -> Result<&mut Storage, RequireError> {
        if self.accounts.contains_key(&address) {
            match self.accounts.get_mut(&address).unwrap() {
                &mut Account::Full {
                    ref mut changing_storage,
                    ..
                } => return Ok(changing_storage),
                &mut Account::Create {
                    ref mut storage,
                    ..
                } => return Ok(storage),
                _ => (),
            }
        }

        return Err(RequireError::Account(address));
    }

    /// Create a new account (that should not yet have existed
    /// before).
    pub fn create(&mut self, address: Address, balance: U256, code: &[u8]) {
        self.accounts.insert(address, Account::Full {
            address, balance, changing_storage: Storage::new(address, false), code: code.into(), nonce: M256::zero(),
        });
    }

    /// Increase the balance of an account.
    pub fn increase_balance(&mut self, address: Address, topup: U256) {
        if topup == U256::zero() { return; }
        let account = match self.accounts.remove(&address) {
            Some(Account::Full {
                address,
                balance,
                changing_storage,
                code,
                nonce,
            }) => {
                Some(Account::Full {
                    address,
                    balance: balance + topup,
                    changing_storage,
                    code,
                    nonce,
                })
            },
            Some(Account::IncreaseBalance(address, balance)) => {
                Some(Account::IncreaseBalance(address, balance + topup))
            },
            Some(Account::DecreaseBalance(address, balance)) => {
                if balance == topup {
                    None
                } else if balance > topup {
                    Some(Account::DecreaseBalance(address, balance - topup))
                } else {
                    Some(Account::IncreaseBalance(address, topup - balance))
                }
            },
            Some(Account::Remove(_)) => {
                panic!()
            },
            Some(Account::Create {
                address,
                balance,
                storage,
                code,
                nonce,
            }) => {
                Some(Account::Create {
                    address,
                    balance: balance + topup,
                    storage,
                    code,
                    nonce,
                })
            },
            None => {
                Some(Account::IncreaseBalance(address, topup))
            },
        };
        if account.is_some() {
            self.accounts.insert(address, account.unwrap());
        }
    }

    /// Decrease the balance of an account.
    pub fn decrease_balance(&mut self, address: Address, withdraw: U256) {
        if withdraw == U256::zero() { return; }
        let account = match self.accounts.remove(&address) {
            Some(Account::Full {
                address,
                balance,
                changing_storage,
                code,
                nonce,
            }) => {
                Some(Account::Full {
                    address,
                    balance: balance - withdraw,
                    changing_storage,
                    code,
                    nonce,
                })
            },
            Some(Account::DecreaseBalance(address, balance)) => {
                Some(Account::DecreaseBalance(address, balance + withdraw))
            },
            Some(Account::IncreaseBalance(address, balance)) => {
                if balance == withdraw {
                    None
                } else if balance > withdraw {
                    Some(Account::IncreaseBalance(address, balance - withdraw))
                } else {
                    Some(Account::DecreaseBalance(address, withdraw - balance))
                }
            },
            Some(Account::Remove(_)) => {
                panic!()
            },
            Some(Account::Create {
                address,
                balance,
                storage,
                code,
                nonce,
            }) => {
                Some(Account::Create {
                    address,
                    balance: balance - withdraw,
                    storage,
                    code,
                    nonce,
                })
            },
            None => {
                Some(Account::DecreaseBalance(address, withdraw))
            },
        };
        if account.is_some() {
            self.accounts.insert(address, account.unwrap());
        }
    }

    /// Set nonce of an account. If the account is not already
    /// commited, returns a `RequireError`.
    pub fn set_nonce(&mut self, address: Address, new_nonce: M256) -> Result<(), RequireError> {
        match self.accounts.get_mut(&address) {
            Some(&mut Account::Full {
                ref mut nonce,
                ..
            }) => {
                *nonce = new_nonce;
                Ok(())
            },
            _ => {
                Err(RequireError::Account(address))
            },
        }
    }

    /// Delete an account from this account state. The account is set
    /// to null. If the account is not already commited, returns a
    /// `RequireError`.
    pub fn remove(&mut self, address: Address) -> Result<(), RequireError> {
        let account = match self.accounts.remove(&address) {
            Some(Account::Full {
                address,
                ..
            }) => {
                Account::Remove(address)
            },
            Some(Account::DecreaseBalance(address, _)) => {
                return Err(RequireError::Account(address));
            },
            Some(Account::IncreaseBalance(address, _)) => {
                return Err(RequireError::Account(address));
            },
            Some(Account::Remove(_)) => {
                panic!()
            },
            Some(Account::Create {
                address,
                ..
            }) => {
                Account::Remove(address)
            },
            None => {
                return Err(RequireError::Account(address));
            },
        };
        self.accounts.insert(address, account);
        Ok(())
    }
}
<commit_msg>Problem: account commitment management for create is incomplete<commit_after>//! Account commitment managment

use std::collections::hash_map::{self, HashMap};
use utils::address::Address;
use utils::bigint::{M256, U256};

use vm::Storage;
use vm::errors::{RequireError, CommitError};

#[derive(Debug, Clone)]
/// A single account commitment.
pub enum AccountCommitment {
    /// Full account commitment. The client that committed account
    /// should not change the account in other EVMs if it decides to
    /// accept the result.
    Full {
        nonce: M256,
        address: Address,
        balance: U256,
        code: Vec<u8>,
    },
    /// Commit only code of the account. The client can keep changing
    /// it in other EVMs if the code remains unchanged.
    Code {
        address: Address,
        code: Vec<u8>,
    },
    /// Commit a storage. Must be used given a full account.
    Storage {
        address: Address,
        index: M256,
        value: M256,
    },
}

impl AccountCommitment {
    /// Address of this account commitment.
    pub fn address(&self) -> Address {
        match self {
            &AccountCommitment::Full {
                address,
                ..
            } => address,
            &AccountCommitment::Code {
                address,
                ..
            } => address,
            &AccountCommitment::Storage {
                address,
                ..
            } => address,
        }
    }
}

#[derive(Debug, Clone)]
/// Represents an account. This is usually returned by the EVM.
pub enum Account {
    /// A full account. The client is expected to replace its own account state with this.
    Full {
        nonce: M256,
        address: Address,
        balance: U256,
        changing_storage: Storage,
        code: Vec<u8>,
    },
    /// Only balance is changed, and it is increasing for this address.
    IncreaseBalance(Address, U256),
    /// Only balance is changed, and it is decreasing for this address.
    DecreaseBalance(Address, U256),
    /// Remove an account, the client is expected to handle only
    /// removal without any ether transfer.
    Remove(Address),
    /// Create a new account.
    Create {
        nonce: M256,
        address: Address,
        balance: U256,
        storage: Storage,
        code: Vec<u8>,
    },
}

impl Account {
    /// Address of this account.
    pub fn address(&self) -> Address {
        match self {
            &Account::Full {
                address,
                ..
            } => address,
            &Account::IncreaseBalance(address, _) => address,
            &Account::DecreaseBalance(address, _) => address,
            &Account::Remove(address) => address,
            &Account::Create {
                address,
                ..
            } => address,
        }
    }
}

#[derive(Debug, Clone)]
/// A struct that manages the current account state for one EVM.
pub struct AccountState {
    accounts: HashMap<Address, Account>,
    codes: HashMap<Address, Vec<u8>>,
}

impl Default for AccountState {
    fn default() -> Self {
        Self {
            accounts: HashMap::new(),
            codes: HashMap::new(),
        }
    }
}

impl AccountState {
    /// Returns all accounts right now in this account state.
    pub fn accounts(&self) -> hash_map::Values<Address, Account> {
        self.accounts.values()
    }

    /// Returns Ok(()) if a full account is in this account
    /// state. Otherwise raise a `RequireError`.
    pub fn require(&self, address: Address) -> Result<(), RequireError> {
        match self.accounts.get(&address) {
            Some(&Account::Full { .. }) => return Ok(()),
            _ => return Err(RequireError::Account(address)),
        }
    }

    /// Returns Ok(()) if either a full account or a partial code
    /// account is in this account state. Otherwise raise a
    /// `RequireError`.
    pub fn require_code(&self, address: Address) -> Result<(), RequireError> {
        if self.codes.contains_key(&address) {
            return Ok(());
        }
        match self.accounts.get(&address) {
            Some(&Account::Full { .. }) => return Ok(()),
            _ => return Err(RequireError::AccountCode(address)),
        }
    }

    /// Returns Ok(()) if the storage exists in the VM. Otherwise
    /// raise a `RequireError`.
    pub fn require_storage(&self, address: Address, index: M256) -> Result<(), RequireError> {
        self.storage(address)?.read(index).and_then(|_| Ok(()))
    }

    /// Commit an account commitment into this account state.
    pub fn commit(&mut self, commitment: AccountCommitment) -> Result<(), CommitError> {
        match commitment {
            AccountCommitment::Full {
                nonce,
                address,
                balance,
                code
            } => {
                let account = if self.accounts.contains_key(&address) {
                    match self.accounts.remove(&address).unwrap() {
                        Account::Full { .. } => return Err(CommitError::AlreadyCommitted),
                        Account::Create { .. } => return Err(CommitError::AlreadyCommitted),
                        Account::Remove(address) => Account::Remove(address),
                        Account::IncreaseBalance(address, topup) => {
                            Account::Full {
                                nonce,
                                address,
                                balance: balance + topup,
                                changing_storage: Storage::new(address, true),
                                code,
                            }
                        },
                        Account::DecreaseBalance(address, withdraw) => {
                            Account::Full {
                                nonce,
                                address,
                                balance: balance - withdraw,
                                changing_storage: Storage::new(address, true),
                                code,
                            }
                        },
                    }
                } else {
                    Account::Full {
                        nonce,
                        address,
                        balance,
                        changing_storage: Storage::new(address, true),
                        code,
                    }
                };

                self.accounts.insert(address, account);
            },
            AccountCommitment::Code {
                address,
                code,
            } => {
                if self.accounts.contains_key(&address) || self.codes.contains_key(&address) {
                    return Err(CommitError::AlreadyCommitted);
                }

                self.codes.insert(address, code);
            },
            AccountCommitment::Storage {
                address,
                index,
                value
            } => {
                match self.accounts.get_mut(&address) {
                    Some(&mut Account::Full {
                        ref mut changing_storage,
                        ..
                    }) => {
                        changing_storage.commit(index, value)?;
                    },
                    _ => {
                        return Err(CommitError::InvalidCommitment);
                    },
                }
            }
        }
        Ok(())
    }

    /// Find code by its address in this account state. If the search
    /// failed, returns a `RequireError`.
    pub fn code(&self, address: Address) -> Result<&[u8], RequireError> {
        if self.codes.contains_key(&address) {
            return Ok(self.codes.get(&address).unwrap().as_slice());
        }

        if self.accounts.contains_key(&address) {
            match self.accounts.get(&address).unwrap() {
                &Account::Full {
                    ref code,
                    ..
                } => return Ok(code.as_slice()),
                &Account::Create {
                    ref code,
                    ..
                } => return Ok(code.as_slice()),
                &Account::Remove(_) => return Ok(&[]),
                &Account::IncreaseBalance(address, _) => return Err(RequireError::Account(address)),
                &Account::DecreaseBalance(address, _) => return Err(RequireError::Account(address)),
            }
        }

        return Err(RequireError::AccountCode(address));
    }

    /// Find nonce by its address in this account state. If the search
    /// failed, returns a `RequireError`.
    pub fn nonce(&self, address: Address) -> Result<M256, RequireError> {
        if self.accounts.contains_key(&address) {
            match self.accounts.get(&address).unwrap() {
                &Account::Full {
                    nonce,
                    ..
                } => return Ok(nonce),
                _ => (),
            }
        }

        return Err(RequireError::Account(address));
    }

    /// Find balance by its address in this account state. If the
    /// search failed, returns a `RequireError`.
    pub fn balance(&self, address: Address) -> Result<U256, RequireError> {
        if self.accounts.contains_key(&address) {
            match self.accounts.get(&address).unwrap() {
                &Account::Full {
                    balance,
                    ..
                } => return Ok(balance),
                _ => (),
            }
        }

        return Err(RequireError::Account(address));
    }

    /// Returns the storage of an account. If the account is not yet
    /// committed, returns a `RequireError`.
    pub fn storage(&self, address: Address) -> Result<&Storage, RequireError> {
        if self.accounts.contains_key(&address) {
            match self.accounts.get(&address).unwrap() {
                &Account::Full {
                    ref changing_storage,
                    ..
                } => return Ok(changing_storage),
                &Account::Create {
                    ref storage,
                    ..
                } => return Ok(storage),
                _ => (),
            }
        }

        return Err(RequireError::Account(address));
    }

    /// Returns the mutable storage of an account. If the account is
    /// not yet committed. returns a `RequireError`.
    pub fn storage_mut(&mut self, address: Address) -> Result<&mut Storage, RequireError> {
        if self.accounts.contains_key(&address) {
            match self.accounts.get_mut(&address).unwrap() {
                &mut Account::Full {
                    ref mut changing_storage,
                    ..
                } => return Ok(changing_storage),
                &mut Account::Create {
                    ref mut storage,
                    ..
                } => return Ok(storage),
                _ => (),
            }
        }

        return Err(RequireError::Account(address));
    }

    /// Create a new account (that should not yet have existed
    /// before).
    pub fn create(&mut self, address: Address, balance: U256, code: &[u8]) {
        let account = if self.accounts.contains_key(&address) {
            match self.accounts.remove(&address).unwrap() {
                Account::Full { .. } => panic!(),
                Account::Create { .. } => panic!(),
                Account::Remove(address) => {
                    Account::Create {
                        address, balance, code: code.into(), nonce: M256::zero(),
                        storage: Storage::new(address, false)
                    }
                },
                Account::IncreaseBalance(address, topup) => {
                    Account::Create {
                        address, code: code.into(), nonce: M256::zero(),
                        balance: balance + topup, storage: Storage::new(address, false)
                    }
                },
                Account::DecreaseBalance(address, withdraw) => {
                    Account::Create {
                        address, code: code.into(), nonce: M256::zero(),
                        balance: balance - withdraw, storage: Storage::new(address, false)
                    }
                },
            }
        } else {
            Account::Create {
                address, balance, code: code.into(), nonce: M256::zero(),
                storage: Storage::new(address, false)
            }
        };

        self.accounts.insert(address, account);
    }

    /// Increase the balance of an account.
    pub fn increase_balance(&mut self, address: Address, topup: U256) {
        if topup == U256::zero() { return; }
        let account = match self.accounts.remove(&address) {
            Some(Account::Full {
                address,
                balance,
                changing_storage,
                code,
                nonce,
            }) => {
                Some(Account::Full {
                    address,
                    balance: balance + topup,
                    changing_storage,
                    code,
                    nonce,
                })
            },
            Some(Account::IncreaseBalance(address, balance)) => {
                Some(Account::IncreaseBalance(address, balance + topup))
            },
            Some(Account::DecreaseBalance(address, balance)) => {
                if balance == topup {
                    None
                } else if balance > topup {
                    Some(Account::DecreaseBalance(address, balance - topup))
                } else {
                    Some(Account::IncreaseBalance(address, topup - balance))
                }
            },
            Some(Account::Remove(_)) => {
                panic!()
            },
            Some(Account::Create {
                address,
                balance,
                storage,
                code,
                nonce,
            }) => {
                Some(Account::Create {
                    address,
                    balance: balance + topup,
                    storage,
                    code,
                    nonce,
                })
            },
            None => {
                Some(Account::IncreaseBalance(address, topup))
            },
        };
        if account.is_some() {
            self.accounts.insert(address, account.unwrap());
        }
    }

    /// Decrease the balance of an account.
    pub fn decrease_balance(&mut self, address: Address, withdraw: U256) {
        if withdraw == U256::zero() { return; }
        let account = match self.accounts.remove(&address) {
            Some(Account::Full {
                address,
                balance,
                changing_storage,
                code,
                nonce,
            }) => {
                Some(Account::Full {
                    address,
                    balance: balance - withdraw,
                    changing_storage,
                    code,
                    nonce,
                })
            },
            Some(Account::DecreaseBalance(address, balance)) => {
                Some(Account::DecreaseBalance(address, balance + withdraw))
            },
            Some(Account::IncreaseBalance(address, balance)) => {
                if balance == withdraw {
                    None
                } else if balance > withdraw {
                    Some(Account::IncreaseBalance(address, balance - withdraw))
                } else {
                    Some(Account::DecreaseBalance(address, withdraw - balance))
                }
            },
            Some(Account::Remove(_)) => {
                panic!()
            },
            Some(Account::Create {
                address,
                balance,
                storage,
                code,
                nonce,
            }) => {
                Some(Account::Create {
                    address,
                    balance: balance - withdraw,
                    storage,
                    code,
                    nonce,
                })
            },
            None => {
                Some(Account::DecreaseBalance(address, withdraw))
            },
        };
        if account.is_some() {
            self.accounts.insert(address, account.unwrap());
        }
    }

    /// Set nonce of an account. If the account is not already
    /// commited, returns a `RequireError`.
    pub fn set_nonce(&mut self, address: Address, new_nonce: M256) -> Result<(), RequireError> {
        match self.accounts.get_mut(&address) {
            Some(&mut Account::Full {
                ref mut nonce,
                ..
            }) => {
                *nonce = new_nonce;
                Ok(())
            },
            _ => {
                Err(RequireError::Account(address))
            },
        }
    }

    /// Delete an account from this account state. The account is set
    /// to null. If the account is not already commited, returns a
    /// `RequireError`.
    pub fn remove(&mut self, address: Address) -> Result<(), RequireError> {
        let account = match self.accounts.remove(&address) {
            Some(Account::Full {
                address,
                ..
            }) => {
                Account::Remove(address)
            },
            Some(Account::DecreaseBalance(address, _)) => {
                return Err(RequireError::Account(address));
            },
            Some(Account::IncreaseBalance(address, _)) => {
                return Err(RequireError::Account(address));
            },
            Some(Account::Remove(_)) => {
                panic!()
            },
            Some(Account::Create {
                address,
                ..
            }) => {
                Account::Remove(address)
            },
            None => {
                return Err(RequireError::Account(address));
            },
        };
        self.accounts.insert(address, account);
        Ok(())
    }
}
<|endoftext|>


<commit_before>
/// Represents an API version of Vulkan.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct Version {
    /// Major version number.
    pub major: u16,
    /// Minor version number.
    pub minor: u16,
    /// Patch version number.
    pub patch: u16,
}

// TODO: implement PartialOrd & Ord

impl Version {
    /// Turns a version number given by Vulkan into a `Version` struct.
    #[inline]
    pub fn from_vulkan_version(value: u32) -> Version {
        Version {
            major: ((value & 0xffc00000) >> 22) as u16,
            minor: ((value & 0x003ff000) >> 12) as u16,
            patch:  (value & 0x00000fff) as u16,
        }
    }

    /// Turns a `Version` into a version number accepted by Vulkan.
    ///
    /// # Panic
    ///
    /// Panicks if the values in the `Version` are out of acceptable range.
    #[inline]
    pub fn into_vulkan_version(&self) -> u32 {
        assert!(self.major <= 0x3ff);
        assert!(self.minor <= 0x3ff);
        assert!(self.patch <= 0xfff);

        (self.major as u32) << 22 | (self.minor as u32) << 12 | (self.patch as u32)
    }
}

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

    #[test]
    fn test() {
        let version = Version { major: 1, minor: 0, patch: 0 };
        assert_eq!(version.into_vulkan_version(), 0x400000);
    }
}
<commit_msg>Implement PartialOrd and Ord for Version<commit_after>use std::cmp::Ordering;

/// Represents an API version of Vulkan.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct Version {
    /// Major version number.
    pub major: u16,
    /// Minor version number.
    pub minor: u16,
    /// Patch version number.
    pub patch: u16,
}

impl PartialOrd for Version {
    #[inline]
    fn partial_cmp(&self, other: &Version) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Version {
    fn cmp(&self, other: &Version) -> Ordering {
        match self.major.cmp(&other.major) {
            Ordering::Equal => (),
            o => return o
        };

        match self.minor.cmp(&other.minor) {
            Ordering::Equal => (),
            o => return o
        };

        self.patch.cmp(&other.patch)
    }
}

impl Version {
    /// Turns a version number given by Vulkan into a `Version` struct.
    #[inline]
    pub fn from_vulkan_version(value: u32) -> Version {
        Version {
            major: ((value & 0xffc00000) >> 22) as u16,
            minor: ((value & 0x003ff000) >> 12) as u16,
            patch:  (value & 0x00000fff) as u16,
        }
    }

    /// Turns a `Version` into a version number accepted by Vulkan.
    ///
    /// # Panic
    ///
    /// Panicks if the values in the `Version` are out of acceptable range.
    #[inline]
    pub fn into_vulkan_version(&self) -> u32 {
        assert!(self.major <= 0x3ff);
        assert!(self.minor <= 0x3ff);
        assert!(self.patch <= 0xfff);

        (self.major as u32) << 22 | (self.minor as u32) << 12 | (self.patch as u32)
    }
}

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

    #[test]
    fn into_vk_version() {
        let version = Version { major: 1, minor: 0, patch: 0 };
        assert_eq!(version.into_vulkan_version(), 0x400000);
    }

    #[test]
    fn greater_major() {
        let v1 = Version { major: 1, minor: 0, patch: 0 };
        let v2 = Version { major: 2, minor: 0, patch: 0 };
        assert!(v2 > v1);
    }

    #[test]
    fn greater_minor() {
        let v1 = Version { major: 1, minor: 1, patch: 0 };
        let v2 = Version { major: 1, minor: 3, patch: 0 };
        assert!(v2 > v1);
    }

    #[test]
    fn greater_patch() {
        let v1 = Version { major: 1, minor: 0, patch: 4 };
        let v2 = Version { major: 1, minor: 0, patch: 5 };
        assert!(v2 > v1);
    }
}
<|endoftext|>


<commit_before>use crate::bindings as br;
use crate::ptr_util::read_into_vec_from_ptr;
use crate::{compiler, spirv, ErrorCode};
use std::ffi::CString;
use std::marker::PhantomData;
use std::ptr;

/// A GLSL target.
#[derive(Debug, Clone)]
pub enum Target {}

pub struct TargetData {
    combined_image_samplers_built: bool,
}

impl spirv::Target for Target {
    type Data = TargetData;
}

#[allow(non_snake_case, non_camel_case_types)]
#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq)]
pub enum Version {
    V1_10,
    V1_20,
    V1_30,
    V1_40,
    V1_50,
    V3_30,
    V4_00,
    V4_10,
    V4_20,
    V4_30,
    V4_40,
    V4_50,
    V4_60,
    V1_00Es,
    V3_00Es,
}

#[derive(Debug, Clone)]
pub struct CompilerVertexOptions {
    pub invert_y: bool,
    pub transform_clip_space: bool,
    pub support_nonzero_base_instance: bool,
}

impl Default for CompilerVertexOptions {
    fn default() -> CompilerVertexOptions {
        CompilerVertexOptions {
            invert_y: false,
            transform_clip_space: false,
            support_nonzero_base_instance: true,
        }
    }
}

#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq)]
#[repr(u8)]
pub enum Precision {
    DontCare,
    Low,
    Medium,
    High,
}

#[derive(Debug, Clone)]
pub struct CompilerFragmentOptions {
    pub default_float_precision: Precision,
    pub default_int_precision: Precision,
}

impl Default for CompilerFragmentOptions {
    fn default() -> CompilerFragmentOptions {
        CompilerFragmentOptions {
            default_float_precision: Precision::Medium,
            default_int_precision: Precision::High,
        }
    }
}

/// GLSL compiler options.
#[derive(Debug, Clone)]
pub struct CompilerOptions {
    pub version: Version,
    pub force_temporary: bool,
    pub vulkan_semantics: bool,
    pub separate_shader_objects: bool,
    pub flatten_multidimensional_arrays: bool,
    pub enable_420_pack_extension: bool,
    pub emit_push_constant_as_uniform_buffer: bool,
    pub emit_uniform_buffer_as_plain_uniforms: bool,
    pub emit_line_directives: bool,
    pub enable_storage_image_qualifier_deduction: bool,
    pub force_zero_initialized_variables: bool,
    pub vertex: CompilerVertexOptions,
    pub fragment: CompilerFragmentOptions,
}

impl CompilerOptions {
    fn as_raw(&self) -> br::ScGlslCompilerOptions {
        use self::Version::*;
        let (version, es) = match self.version {
            V1_10 => (1_10, false),
            V1_20 => (1_20, false),
            V1_30 => (1_30, false),
            V1_40 => (1_40, false),
            V1_50 => (1_50, false),
            V3_30 => (3_30, false),
            V4_00 => (4_00, false),
            V4_10 => (4_10, false),
            V4_20 => (4_20, false),
            V4_30 => (4_30, false),
            V4_40 => (4_40, false),
            V4_50 => (4_50, false),
            V4_60 => (4_60, false),
            V1_00Es => (1_00, true),
            V3_00Es => (3_00, true),
        };
        br::ScGlslCompilerOptions {
            vertex_invert_y: self.vertex.invert_y,
            vertex_transform_clip_space: self.vertex.transform_clip_space,
            version,
            es,
            vertex_support_nonzero_base_instance: self.vertex.support_nonzero_base_instance,
            fragment_default_float_precision: self.fragment.default_float_precision as u8,
            fragment_default_int_precision: self.fragment.default_int_precision as u8,
            force_temporary: self.force_temporary,
            vulkan_semantics: self.vulkan_semantics,
            separate_shader_objects: self.separate_shader_objects,
            flatten_multidimensional_arrays: self.flatten_multidimensional_arrays,
            enable_420_pack_extension: self.enable_420_pack_extension,
            emit_push_constant_as_uniform_buffer: self.emit_push_constant_as_uniform_buffer,
            emit_uniform_buffer_as_plain_uniforms: self.emit_uniform_buffer_as_plain_uniforms,
            emit_line_directives: self.emit_line_directives,
            enable_storage_image_qualifier_deduction: self.enable_storage_image_qualifier_deduction,
            force_zero_initialized_variables: self.force_zero_initialized_variables,
        }
    }
}

impl Default for CompilerOptions {
    fn default() -> CompilerOptions {
        CompilerOptions {
            version: Version::V4_50,
            force_temporary: false,
            vulkan_semantics: false,
            separate_shader_objects: false,
            flatten_multidimensional_arrays: false,
            enable_420_pack_extension: true,
            emit_push_constant_as_uniform_buffer: false,
            emit_uniform_buffer_as_plain_uniforms: false,
            emit_line_directives: false,
            enable_storage_image_qualifier_deduction: true,
            force_zero_initialized_variables: false,
            vertex: CompilerVertexOptions::default(),
            fragment: CompilerFragmentOptions::default(),
        }
    }
}

impl spirv::Parse<Target> for spirv::Ast<Target> {
    fn parse(module: &spirv::Module) -> Result<Self, ErrorCode> {
        let compiler = {
            let mut compiler = ptr::null_mut();
            unsafe {
                check!(br::sc_internal_compiler_glsl_new(
                    &mut compiler,
                    module.words.as_ptr() as *const u32,
                    module.words.len() as usize,
                ));
            }

            compiler::Compiler {
                sc_compiler: compiler,
                target_data: TargetData {
                    combined_image_samplers_built: false,
                },
                has_been_compiled: false,
            }
        };

        Ok(spirv::Ast {
            compiler,
            target_type: PhantomData,
        })
    }
}

impl spirv::Compile<Target> for spirv::Ast<Target> {
    type CompilerOptions = CompilerOptions;

    /// Set GLSL compiler specific compilation settings.
    fn set_compiler_options(&mut self, options: &CompilerOptions) -> Result<(), ErrorCode> {
        let raw_options = options.as_raw();
        unsafe {
            check!(br::sc_internal_compiler_glsl_set_options(
                self.compiler.sc_compiler,
                &raw_options,
            ));
        }

        Ok(())
    }

    /// Generate GLSL shader from the AST.
    fn compile(&mut self) -> Result<String, ErrorCode> {
        self.build_combined_image_samplers()?;
        self.compiler.compile()
    }
}

impl spirv::Ast<Target> {
    pub fn build_combined_image_samplers(&mut self) -> Result<(), ErrorCode> {
        unsafe {
            if !self.compiler.target_data.combined_image_samplers_built {
                check!(br::sc_internal_compiler_glsl_build_combined_image_samplers(
                    self.compiler.sc_compiler
                ));
                self.compiler.target_data.combined_image_samplers_built = true
            }
        }

        Ok(())
    }

    pub fn get_combined_image_samplers(
        &mut self,
    ) -> Result<Vec<spirv::CombinedImageSampler>, ErrorCode> {
        self.build_combined_image_samplers()?;
        unsafe {
            let mut samplers_raw: *const br::ScCombinedImageSampler = std::ptr::null();
            let mut samplers_raw_length: usize = 0;

            check!(br::sc_internal_compiler_glsl_get_combined_image_samplers(
                self.compiler.sc_compiler,
                &mut samplers_raw as _,
                &mut samplers_raw_length as _,
            ));

            let samplers = read_into_vec_from_ptr(samplers_raw, samplers_raw_length)
                .iter()
                .map(|sc| spirv::CombinedImageSampler {
                    combined_id: sc.combined_id,
                    image_id: sc.image_id,
                    sampler_id: sc.sampler_id,
                })
                .collect();

            Ok(samplers)
        }
    }

    pub fn add_header_line(&mut self, line: &str) -> Result<(), ErrorCode> {
        unsafe {
            let line = CString::new(line);
            match line {
                Ok(line) => {
                    check!(br::sc_internal_compiler_glsl_add_header_line(
                        self.compiler.sc_compiler,
                        line.as_ptr(),
                    ));
                }
                _ => return Err(ErrorCode::Unhandled),
            }

            Ok(())
        }
    }

    pub fn flatten_buffer_block(&mut self, id: u32) -> Result<(), ErrorCode> {
        unsafe {
            check!(br::sc_internal_compiler_glsl_flatten_buffer_block(
                self.compiler.sc_compiler,
                id,
            ));

            Ok(())
        }
    }
}
<commit_msg>Explicitly specify variable values for Precision<commit_after>use crate::bindings as br;
use crate::ptr_util::read_into_vec_from_ptr;
use crate::{compiler, spirv, ErrorCode};
use std::ffi::CString;
use std::marker::PhantomData;
use std::ptr;

/// A GLSL target.
#[derive(Debug, Clone)]
pub enum Target {}

pub struct TargetData {
    combined_image_samplers_built: bool,
}

impl spirv::Target for Target {
    type Data = TargetData;
}

#[allow(non_snake_case, non_camel_case_types)]
#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq)]
pub enum Version {
    V1_10,
    V1_20,
    V1_30,
    V1_40,
    V1_50,
    V3_30,
    V4_00,
    V4_10,
    V4_20,
    V4_30,
    V4_40,
    V4_50,
    V4_60,
    V1_00Es,
    V3_00Es,
}

#[derive(Debug, Clone)]
pub struct CompilerVertexOptions {
    pub invert_y: bool,
    pub transform_clip_space: bool,
    pub support_nonzero_base_instance: bool,
}

impl Default for CompilerVertexOptions {
    fn default() -> CompilerVertexOptions {
        CompilerVertexOptions {
            invert_y: false,
            transform_clip_space: false,
            support_nonzero_base_instance: true,
        }
    }
}

// Note: These values should match with `CompilerGLSL::Options::Precision`.
#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq)]
#[repr(u8)]
pub enum Precision {
    DontCare = 0,
    Low = 1,
    Medium = 2,
    High = 3,
}

#[derive(Debug, Clone)]
pub struct CompilerFragmentOptions {
    pub default_float_precision: Precision,
    pub default_int_precision: Precision,
}

impl Default for CompilerFragmentOptions {
    fn default() -> CompilerFragmentOptions {
        CompilerFragmentOptions {
            default_float_precision: Precision::Medium,
            default_int_precision: Precision::High,
        }
    }
}

/// GLSL compiler options.
#[derive(Debug, Clone)]
pub struct CompilerOptions {
    pub version: Version,
    pub force_temporary: bool,
    pub vulkan_semantics: bool,
    pub separate_shader_objects: bool,
    pub flatten_multidimensional_arrays: bool,
    pub enable_420_pack_extension: bool,
    pub emit_push_constant_as_uniform_buffer: bool,
    pub emit_uniform_buffer_as_plain_uniforms: bool,
    pub emit_line_directives: bool,
    pub enable_storage_image_qualifier_deduction: bool,
    pub force_zero_initialized_variables: bool,
    pub vertex: CompilerVertexOptions,
    pub fragment: CompilerFragmentOptions,
}

impl CompilerOptions {
    fn as_raw(&self) -> br::ScGlslCompilerOptions {
        use self::Version::*;
        let (version, es) = match self.version {
            V1_10 => (1_10, false),
            V1_20 => (1_20, false),
            V1_30 => (1_30, false),
            V1_40 => (1_40, false),
            V1_50 => (1_50, false),
            V3_30 => (3_30, false),
            V4_00 => (4_00, false),
            V4_10 => (4_10, false),
            V4_20 => (4_20, false),
            V4_30 => (4_30, false),
            V4_40 => (4_40, false),
            V4_50 => (4_50, false),
            V4_60 => (4_60, false),
            V1_00Es => (1_00, true),
            V3_00Es => (3_00, true),
        };
        br::ScGlslCompilerOptions {
            vertex_invert_y: self.vertex.invert_y,
            vertex_transform_clip_space: self.vertex.transform_clip_space,
            version,
            es,
            vertex_support_nonzero_base_instance: self.vertex.support_nonzero_base_instance,
            fragment_default_float_precision: self.fragment.default_float_precision as u8,
            fragment_default_int_precision: self.fragment.default_int_precision as u8,
            force_temporary: self.force_temporary,
            vulkan_semantics: self.vulkan_semantics,
            separate_shader_objects: self.separate_shader_objects,
            flatten_multidimensional_arrays: self.flatten_multidimensional_arrays,
            enable_420_pack_extension: self.enable_420_pack_extension,
            emit_push_constant_as_uniform_buffer: self.emit_push_constant_as_uniform_buffer,
            emit_uniform_buffer_as_plain_uniforms: self.emit_uniform_buffer_as_plain_uniforms,
            emit_line_directives: self.emit_line_directives,
            enable_storage_image_qualifier_deduction: self.enable_storage_image_qualifier_deduction,
            force_zero_initialized_variables: self.force_zero_initialized_variables,
        }
    }
}

impl Default for CompilerOptions {
    fn default() -> CompilerOptions {
        CompilerOptions {
            version: Version::V4_50,
            force_temporary: false,
            vulkan_semantics: false,
            separate_shader_objects: false,
            flatten_multidimensional_arrays: false,
            enable_420_pack_extension: true,
            emit_push_constant_as_uniform_buffer: false,
            emit_uniform_buffer_as_plain_uniforms: false,
            emit_line_directives: false,
            enable_storage_image_qualifier_deduction: true,
            force_zero_initialized_variables: false,
            vertex: CompilerVertexOptions::default(),
            fragment: CompilerFragmentOptions::default(),
        }
    }
}

impl spirv::Parse<Target> for spirv::Ast<Target> {
    fn parse(module: &spirv::Module) -> Result<Self, ErrorCode> {
        let compiler = {
            let mut compiler = ptr::null_mut();
            unsafe {
                check!(br::sc_internal_compiler_glsl_new(
                    &mut compiler,
                    module.words.as_ptr() as *const u32,
                    module.words.len() as usize,
                ));
            }

            compiler::Compiler {
                sc_compiler: compiler,
                target_data: TargetData {
                    combined_image_samplers_built: false,
                },
                has_been_compiled: false,
            }
        };

        Ok(spirv::Ast {
            compiler,
            target_type: PhantomData,
        })
    }
}

impl spirv::Compile<Target> for spirv::Ast<Target> {
    type CompilerOptions = CompilerOptions;

    /// Set GLSL compiler specific compilation settings.
    fn set_compiler_options(&mut self, options: &CompilerOptions) -> Result<(), ErrorCode> {
        let raw_options = options.as_raw();
        unsafe {
            check!(br::sc_internal_compiler_glsl_set_options(
                self.compiler.sc_compiler,
                &raw_options,
            ));
        }

        Ok(())
    }

    /// Generate GLSL shader from the AST.
    fn compile(&mut self) -> Result<String, ErrorCode> {
        self.build_combined_image_samplers()?;
        self.compiler.compile()
    }
}

impl spirv::Ast<Target> {
    pub fn build_combined_image_samplers(&mut self) -> Result<(), ErrorCode> {
        unsafe {
            if !self.compiler.target_data.combined_image_samplers_built {
                check!(br::sc_internal_compiler_glsl_build_combined_image_samplers(
                    self.compiler.sc_compiler
                ));
                self.compiler.target_data.combined_image_samplers_built = true
            }
        }

        Ok(())
    }

    pub fn get_combined_image_samplers(
        &mut self,
    ) -> Result<Vec<spirv::CombinedImageSampler>, ErrorCode> {
        self.build_combined_image_samplers()?;
        unsafe {
            let mut samplers_raw: *const br::ScCombinedImageSampler = std::ptr::null();
            let mut samplers_raw_length: usize = 0;

            check!(br::sc_internal_compiler_glsl_get_combined_image_samplers(
                self.compiler.sc_compiler,
                &mut samplers_raw as _,
                &mut samplers_raw_length as _,
            ));

            let samplers = read_into_vec_from_ptr(samplers_raw, samplers_raw_length)
                .iter()
                .map(|sc| spirv::CombinedImageSampler {
                    combined_id: sc.combined_id,
                    image_id: sc.image_id,
                    sampler_id: sc.sampler_id,
                })
                .collect();

            Ok(samplers)
        }
    }

    pub fn add_header_line(&mut self, line: &str) -> Result<(), ErrorCode> {
        unsafe {
            let line = CString::new(line);
            match line {
                Ok(line) => {
                    check!(br::sc_internal_compiler_glsl_add_header_line(
                        self.compiler.sc_compiler,
                        line.as_ptr(),
                    ));
                }
                _ => return Err(ErrorCode::Unhandled),
            }

            Ok(())
        }
    }

    pub fn flatten_buffer_block(&mut self, id: u32) -> Result<(), ErrorCode> {
        unsafe {
            check!(br::sc_internal_compiler_glsl_flatten_buffer_block(
                self.compiler.sc_compiler,
                id,
            ));

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


<commit_before>/*!
Hypervisor API for OSX

To use this library, you need

* OS X Yosemite (10.10) or newer

* an Intel processor with the VT-x feature set that includes Extended Page
Tables (EPT) and Unrestricted Mode. To verify this, run and expect the following
in your Terminal:

  ```shell
  $ sysctl kern.hv_support
  kern.hv_support: 1
  ```
!*/
//extern crate core;

use super::core::fmt;
use libc::*;

#[allow(non_camel_case_types)]
pub mod ffi;

use self::ffi::*;
use self::ffi::consts::*;

/// Error returned after every call
pub enum Error {
    /// Success
    Success,
    /// Error
    Error,
    /// Busy
    Busy,
    /// Bad argument
    BadArg,
    /// No resources
    NoRes,
    /// No device
    NoDev,
    /// Unsupported
    Unsupp
}

impl fmt::Debug for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::Success => write!(f, "Success"),
            Error::Error => write!(f, "Error"),
            Error::Busy => write!(f, "Busy"),
            Error::BadArg => write!(f, "Bad argument"),
            Error::NoRes => write!(f, "No resources"),
            Error::NoDev => write!(f, "No device"),
            Error::Unsupp => write!(f, "Unsupported"),
        }
    }
}

fn match_error_code(code: hv_return_t) -> Error {
    match code {
        HV_SUCCESS => Error::Success,
        HV_BUSY => Error::Busy,
        HV_BAD_ARGUMENT => Error::BadArg,
        HV_NO_RESOURCES => Error::NoRes,
        HV_NO_DEVICE => Error::NoDev,
        HV_UNSUPPORTED => Error::Unsupp,
        _ => Error::Error
    }
}

/// Creates a VM instance for the current task
pub fn create_vm() -> Error {
    match_error_code(unsafe {
        hv_vm_create(0)
    })
}

/// Destroys the VM instance associated with the current task
pub fn destory_vm() -> Error {
    match_error_code(unsafe {
        hv_vm_destroy()
    })
}

/// Virtual CPU
#[allow(non_camel_case_types)]
pub struct vCPU {
    /// Virtual CPU ID
    pub id: u32
}

impl vCPU {
    /// Creates a vCPU instance for the current thread
    pub fn new() -> Result<vCPU, Error> {
        let mut vcpuid: hv_vcpuid_t = 0;

        let error = match_error_code(unsafe {
            hv_vcpu_create(&mut vcpuid, 0)
        });

        match error {
            Error::Success => Ok(vCPU {
                id: vcpuid as u32
            }),
            _ => Err(error)
        }
    }

    /// Destroys the vCPU instance associated with the current thread
    pub fn destory(&self) -> Error {
        match_error_code(unsafe {
            hv_vcpu_destroy(self.id as hv_vcpuid_t)
        })
    }

    /// Executes the vCPU
    pub fn run(&self) -> Error {
        match_error_code(unsafe {
            hv_vcpu_run(self.id as hv_vcpuid_t)
        })
    }

    /// Returns the cumulative execution time of the vCPU in nanoseconds
    pub fn exec_time(&self) -> Result<u64, Error> {
        let mut exec_time: uint64_t = 0;

        let error = match_error_code(unsafe {
            hv_vcpu_get_exec_time(self.id, &mut exec_time)
        });

        match error {
            Error::Success => Ok(exec_time as u64),
            _ => Err(error)
        }
    }

}

impl fmt::Debug for vCPU {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "vCPU ID: {}", (*self).id)
    }
}
<commit_msg>Add two more functions to manage vCPU state flush and invalidate_tlb<commit_after>/*!
Hypervisor API for OSX

To use this library, you need

* OS X Yosemite (10.10) or newer

* an Intel processor with the VT-x feature set that includes Extended Page
Tables (EPT) and Unrestricted Mode. To verify this, run and expect the following
in your Terminal:

  ```shell
  $ sysctl kern.hv_support
  kern.hv_support: 1
  ```
!*/
//extern crate core;

use super::core::fmt;
use libc::*;

#[allow(non_camel_case_types)]
pub mod ffi;

use self::ffi::*;
use self::ffi::consts::*;

/// Error returned after every call
pub enum Error {
    /// Success
    Success,
    /// Error
    Error,
    /// Busy
    Busy,
    /// Bad argument
    BadArg,
    /// No resources
    NoRes,
    /// No device
    NoDev,
    /// Unsupported
    Unsupp
}

impl fmt::Debug for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::Success => write!(f, "Success"),
            Error::Error => write!(f, "Error"),
            Error::Busy => write!(f, "Busy"),
            Error::BadArg => write!(f, "Bad argument"),
            Error::NoRes => write!(f, "No resources"),
            Error::NoDev => write!(f, "No device"),
            Error::Unsupp => write!(f, "Unsupported"),
        }
    }
}

fn match_error_code(code: hv_return_t) -> Error {
    match code {
        HV_SUCCESS => Error::Success,
        HV_BUSY => Error::Busy,
        HV_BAD_ARGUMENT => Error::BadArg,
        HV_NO_RESOURCES => Error::NoRes,
        HV_NO_DEVICE => Error::NoDev,
        HV_UNSUPPORTED => Error::Unsupp,
        _ => Error::Error
    }
}

/// Creates a VM instance for the current task
pub fn create_vm() -> Error {
    match_error_code(unsafe {
        hv_vm_create(0)
    })
}

/// Destroys the VM instance associated with the current task
pub fn destory_vm() -> Error {
    match_error_code(unsafe {
        hv_vm_destroy()
    })
}

/// Virtual CPU
#[allow(non_camel_case_types)]
pub struct vCPU {
    /// Virtual CPU ID
    pub id: u32
}

impl vCPU {
    /// Creates a vCPU instance for the current thread
    pub fn new() -> Result<vCPU, Error> {
        let mut vcpuid: hv_vcpuid_t = 0;

        let error = match_error_code(unsafe {
            hv_vcpu_create(&mut vcpuid, 0)
        });

        match error {
            Error::Success => Ok(vCPU {
                id: vcpuid as u32
            }),
            _ => Err(error)
        }
    }

    /// Destroys the vCPU instance associated with the current thread
    pub fn destory(&self) -> Error {
        match_error_code(unsafe {
            hv_vcpu_destroy(self.id as hv_vcpuid_t)
        })
    }

    /// Executes the vCPU
    pub fn run(&self) -> Error {
        match_error_code(unsafe {
            hv_vcpu_run(self.id as hv_vcpuid_t)
        })
    }

    /// Returns the cumulative execution time of the vCPU in nanoseconds
    pub fn exec_time(&self) -> Result<u64, Error> {
        let mut exec_time: uint64_t = 0;

        let error = match_error_code(unsafe {
            hv_vcpu_get_exec_time(self.id, &mut exec_time)
        });

        match error {
            Error::Success => Ok(exec_time as u64),
            _ => Err(error)
        }
    }

    /// Forces flushing of cached vCPU state
    pub fn flush(&self) -> Error {
        match_error_code(unsafe {
            hv_vcpu_flush(self.id as hv_vcpuid_t)
        })
    }

    /// Invalidates the translation lookaside buffer (TLB) of a vCPU
    pub fn invalidate_tlb(&self) -> Error {
        match_error_code(unsafe {
            hv_vcpu_invalidate_tlb(self.id as hv_vcpuid_t)
        })
    }

}

impl fmt::Debug for vCPU {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "vCPU ID: {}", (*self).id)
    }
}
<|endoftext|>


<commit_before>//! A palette for use with Objects

use serde;
use rustc_serialize; // needed for toml
use toml;

use std::fs::File;
use std::path::Path;
use std::error::Error;
use std::io::Read;

use std::collections::HashMap;
use std::iter::{repeat,FromIterator};

use kiss3d::scene::SceneNode;

use objects::ObjectID;

/// An RGB color
#[derive(Eq,PartialEq,Ord,PartialOrd,Hash,Copy,Clone)]
pub struct Color(pub u8,pub u8,pub u8);

pub static DEFAULT_COLORS : [(u8, u8, u8); 11] = [
    ( 77, 175,  74), // Green
    (152,  78, 163), // Purple
    (255, 127,   0), // Orange
    (228,  26,  28), // Red
    ( 55, 126, 184), // Blue
    (166,  86,  40), // Brown
    (247, 129, 191), // Pink
    (153, 153, 153), // Gray
    (255, 255,  51), // Yellow
    (255, 255, 255), // White
    (  0,   0,   0), // Black
];

impl Color {
    fn to_floats(self) -> (f32, f32, f32) {
        let Color(r,g,b) = self;
        (
            (r as f32) / (u8::max_value() as f32),
            (g as f32) / (u8::max_value() as f32),
            (b as f32) / (u8::max_value() as f32),
        )
    }
}

/// A [bool] for keeping track of which part of the objectID should be used
#[derive(Debug,Eq,PartialEq,Ord,PartialOrd,Hash,Clone)]
#[derive(Serialize,Deserialize)]
#[derive(RustcDecodable,RustcEncodable)]
pub struct PartialIDer {
    /// Which sections of an ObjectID should be converted
    pub bools : Vec<bool>,
}

impl PartialIDer {
    /// New, with all set to given boolean value
    pub fn new(n : usize, value : bool) -> Self {
        PartialIDer {
            bools : repeat(value).take(n).collect(),
        }
    }

    /// Get the portion of an ID
    pub fn partial<'a>(&self, name: &'a ObjectID) -> Vec<&'a str> {
        let &ObjectID(ref names) = name;
        names.iter()
            .zip(self.bools.iter())
            .filter_map(|(n, &b)| if b {Some(&n[..])} else {None})
            .collect()
    }

    /// Get the portion of an ID
    pub fn as_id(&self, name: &ObjectID) -> ObjectID {
        let &ObjectID(ref names) = name;
        ObjectID(
            names
            .iter()
            .zip(self.bools.iter())
            .filter_map(|(n, &b)| if b {Some(n.clone())} else {None})
            .collect()
        )
    }

    /// Get a string representing the current state
    pub fn as_string(&self) -> String {
        let mut s = String::with_capacity(self.bools.len());
        let pieces : Vec<String> = self.bools.iter().enumerate().map(
            |(n, &b)| if b {format!("{}", (n+1) % 10)} else {String::from("_")}
        ).collect();
        s.extend(pieces.iter().map(|p| &p[0..1]));
        s
    }
}

/// A way to convert string names to colors
#[derive(Clone)]
pub struct Palette {
    default_colors : Vec<Color>,
    partials : PartialIDer,

    /// Mapping of names -> colors, when found. names not found will be given default colors.
    pub assigned : HashMap<ObjectID, Color>,
    next_color : usize
}

impl Palette {
    /// Set the color of an object to match what the palette says
    pub fn set_color(&mut self, name: &ObjectID, node: &mut SceneNode) {
        let (r, g, b) = self.get_color(name).to_floats();
        node.set_color(r, g, b);
    }

    /// Get the color for a particular ID (using partials mask)
    pub fn get_color(&mut self, name: &ObjectID) -> Color {
        let partial = self.partials.as_id(name);
        let (next_color, default_colors, assigned) =
            (&mut self.next_color, &self.default_colors, &mut self.assigned);
        let color = *assigned.entry(partial).or_insert_with(|| {
            let col = default_colors[*next_color];
            *next_color = (*next_color + 1) % default_colors.len();
            col
        });
        return color
    }

    /// Toggle whether or not to use a certain partial value
    pub fn toggle_partial(&mut self, n : usize){
        self.partials.bools[n] = !self.partials.bools[n];
        self.next_color = 0;
    }

    /// Check value of a certain partial
    pub fn get_partial(&mut self, n : usize) -> bool {
        self.partials.bools[n]
    }

    /// Check value of a certain partial
    pub fn set_partial(&mut self, n : usize, value : bool){
        self.partials.bools[n] = value;
        self.next_color = 0;
    }

    /// Check value of a certain partial
    pub fn set_all_partial(&mut self, value : bool){
        let n = self.partials.bools.len();
        for i in 0..n {
            self.partials.bools[i] = value;
        }
        self.next_color = 0;
    }

    /// Get a string like '_23____' for which parts of the ObjectID we are coloring with respect to
    pub fn partials_string(&self) -> String {
        self.partials.as_string()
    }

    /// Load from a file, using toml-rs and serialize
    pub fn load(path : &Path) -> Result<Self, Box<Error>> {
        let mut file : File = try!(File::open(path));
        let mut s = String::new();
        let _ = try!(file.read_to_string(&mut s));
        let mut parser = toml::Parser::new(s.as_ref());

        let parsed = parser.parse();
        let values = match parsed {
            Some(v) => v,
            None => {
                // We can unwrap here, becase parser.parse() == None means there were errors
                let err : Box<Error> = parser.errors.pop().unwrap().into();
                return Err(err)
            }
        };
        let table : toml::Value = toml::Value::Table(values);
        let mut decoder = toml::Decoder::new(table);
        let p = try!(rustc_serialize::Decodable::decode(&mut decoder));
        Ok(p)
    }
}

impl Default for Palette {
    fn default() -> Self {
        Palette {
            default_colors : FromIterator::from_iter(DEFAULT_COLORS.iter().map(
                |&(r,g,b)|{Color(r,g,b)}
            )),
            partials : PartialIDer::new(8, true),
            assigned : HashMap::new(),
            next_color : 0
        }
    }
}

#[derive(Serialize, RustcEncodable, Copy, Clone)]
struct AssignmentRef<'a> {
    names : &'a ObjectID,
    color : &'a Color
}

fn to_assignments<'a>(assigned : &'a HashMap<ObjectID, Color>) -> Vec<AssignmentRef<'a>> {
    assigned.into_iter().map(|(k, v)| {
        AssignmentRef{
            names : k,
            color : v
        }
    }).collect()
}

/// A reference to a palette, with minimal data, for serialization purposes
#[derive(Serialize, RustcEncodable, Clone)]
pub struct PaletteRef<'a> {
    defaults : &'a Vec<Color>,
    partials : &'a PartialIDer,
    assigned : Vec<AssignmentRef<'a>>
}

impl<'a> From<&'a Palette> for PaletteRef<'a> {
    fn from(palette: &'a Palette) -> Self {
        PaletteRef {
            defaults : &palette.default_colors,
            partials : &palette.partials,
            assigned : to_assignments(&palette.assigned),
        }
    }
}

#[derive(Deserialize, RustcDecodable, Clone)]
struct Assignment {
    names : ObjectID,
    color : Color
}

fn from_assignments(assigned : Vec<Assignment>) -> HashMap<ObjectID, Color> {
    assigned.into_iter().map(|a| {
        (a.names, a.color)
    }).collect()
}

/// A reference to a palette, with minimal data
#[derive(Deserialize, RustcDecodable, Clone)]
pub struct PaletteBasic {
    defaults : Vec<Color>,
    partials : PartialIDer,
    assigned : Option<Vec<Assignment>>,
    next_color : Option<usize>
}

impl From<PaletteBasic> for Palette {
    fn from(p : PaletteBasic) -> Self {
        Palette {
            default_colors : p.defaults,
            partials : p.partials,
            assigned : p.assigned.map(from_assignments).unwrap_or_default(),
            next_color : p.next_color.unwrap_or(0)
        }
    }
}

//TODO: add test for round-trip through toml and json

impl serde::Serialize for Color {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error> where S: serde::Serializer {
        let &Color(r,g,b) = self;
        serde::Serialize::serialize(&(r,g,b), serializer)
    }
}

impl serde::Deserialize for Color {
    fn deserialize<D : serde::Deserializer>(deserializer: &mut D) -> Result<Self, D::Error> {
        let (r,g,b) = try!(serde::Deserialize::deserialize(deserializer));
        Ok(Color(r,g,b))
    }
}

impl rustc_serialize::Encodable for Color {
    fn encode<S: rustc_serialize::Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        let &Color(r,g,b) = self;
        rustc_serialize::Encodable::encode(&(r,g,b), s)
    }
}

impl rustc_serialize::Decodable for Color {
    fn decode<D: rustc_serialize::Decoder>(d: &mut D) -> Result<Self, D::Error> {
        let (r,g,b) = try!(rustc_serialize::Decodable::decode(d));
        Ok(Color(r,g,b))
    }
}

impl serde::Serialize for Palette {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error> where S: serde::Serializer {
        PaletteRef::from(self).serialize(serializer)
    }
}

impl serde::Deserialize for Palette {
    fn deserialize<D : serde::Deserializer>(deserializer: &mut D) -> Result<Self, D::Error>  {
        let p = try!(PaletteBasic::deserialize(deserializer));
        Ok(Palette::from(p))
    }
}

impl rustc_serialize::Encodable for Palette {
    fn encode<S: rustc_serialize::Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        PaletteRef::from(self).encode(s)
    }
}

impl rustc_serialize::Decodable for Palette {
    fn decode<D: rustc_serialize::Decoder>(d: &mut D) -> Result<Self, D::Error> {
        let p = try!(PaletteBasic::decode(d));
        Ok(Palette::from(p))
    }
}
<commit_msg>docs: default colors<commit_after>//! A palette for use with Objects

use serde;
use rustc_serialize; // needed for toml
use toml;

use std::fs::File;
use std::path::Path;
use std::error::Error;
use std::io::Read;

use std::collections::HashMap;
use std::iter::{repeat,FromIterator};

use kiss3d::scene::SceneNode;

use objects::ObjectID;

/// An RGB color
#[derive(Eq,PartialEq,Ord,PartialOrd,Hash,Copy,Clone)]
pub struct Color(pub u8,pub u8,pub u8);

/// Default colors to use, when no others are specified.
pub static DEFAULT_COLORS : [(u8, u8, u8); 11] = [
    ( 77, 175,  74), // Green
    (152,  78, 163), // Purple
    (255, 127,   0), // Orange
    (228,  26,  28), // Red
    ( 55, 126, 184), // Blue
    (166,  86,  40), // Brown
    (247, 129, 191), // Pink
    (153, 153, 153), // Gray
    (255, 255,  51), // Yellow
    (255, 255, 255), // White
    (  0,   0,   0), // Black
];

impl Color {
    fn to_floats(self) -> (f32, f32, f32) {
        let Color(r,g,b) = self;
        (
            (r as f32) / (u8::max_value() as f32),
            (g as f32) / (u8::max_value() as f32),
            (b as f32) / (u8::max_value() as f32),
        )
    }
}

/// A [bool] for keeping track of which part of the objectID should be used
#[derive(Debug,Eq,PartialEq,Ord,PartialOrd,Hash,Clone)]
#[derive(Serialize,Deserialize)]
#[derive(RustcDecodable,RustcEncodable)]
pub struct PartialIDer {
    /// Which sections of an ObjectID should be converted
    pub bools : Vec<bool>,
}

impl PartialIDer {
    /// New, with all set to given boolean value
    pub fn new(n : usize, value : bool) -> Self {
        PartialIDer {
            bools : repeat(value).take(n).collect(),
        }
    }

    /// Get the portion of an ID
    pub fn partial<'a>(&self, name: &'a ObjectID) -> Vec<&'a str> {
        let &ObjectID(ref names) = name;
        names.iter()
            .zip(self.bools.iter())
            .filter_map(|(n, &b)| if b {Some(&n[..])} else {None})
            .collect()
    }

    /// Get the portion of an ID
    pub fn as_id(&self, name: &ObjectID) -> ObjectID {
        let &ObjectID(ref names) = name;
        ObjectID(
            names
            .iter()
            .zip(self.bools.iter())
            .filter_map(|(n, &b)| if b {Some(n.clone())} else {None})
            .collect()
        )
    }

    /// Get a string representing the current state
    pub fn as_string(&self) -> String {
        let mut s = String::with_capacity(self.bools.len());
        let pieces : Vec<String> = self.bools.iter().enumerate().map(
            |(n, &b)| if b {format!("{}", (n+1) % 10)} else {String::from("_")}
        ).collect();
        s.extend(pieces.iter().map(|p| &p[0..1]));
        s
    }
}

/// A way to convert string names to colors
#[derive(Clone)]
pub struct Palette {
    default_colors : Vec<Color>,
    partials : PartialIDer,

    /// Mapping of names -> colors, when found. names not found will be given default colors.
    pub assigned : HashMap<ObjectID, Color>,
    next_color : usize
}

impl Palette {
    /// Set the color of an object to match what the palette says
    pub fn set_color(&mut self, name: &ObjectID, node: &mut SceneNode) {
        let (r, g, b) = self.get_color(name).to_floats();
        node.set_color(r, g, b);
    }

    /// Get the color for a particular ID (using partials mask)
    pub fn get_color(&mut self, name: &ObjectID) -> Color {
        let partial = self.partials.as_id(name);
        let (next_color, default_colors, assigned) =
            (&mut self.next_color, &self.default_colors, &mut self.assigned);
        let color = *assigned.entry(partial).or_insert_with(|| {
            let col = default_colors[*next_color];
            *next_color = (*next_color + 1) % default_colors.len();
            col
        });
        return color
    }

    /// Toggle whether or not to use a certain partial value
    pub fn toggle_partial(&mut self, n : usize){
        self.partials.bools[n] = !self.partials.bools[n];
        self.next_color = 0;
    }

    /// Check value of a certain partial
    pub fn get_partial(&mut self, n : usize) -> bool {
        self.partials.bools[n]
    }

    /// Check value of a certain partial
    pub fn set_partial(&mut self, n : usize, value : bool){
        self.partials.bools[n] = value;
        self.next_color = 0;
    }

    /// Check value of a certain partial
    pub fn set_all_partial(&mut self, value : bool){
        let n = self.partials.bools.len();
        for i in 0..n {
            self.partials.bools[i] = value;
        }
        self.next_color = 0;
    }

    /// Get a string like '_23____' for which parts of the ObjectID we are coloring with respect to
    pub fn partials_string(&self) -> String {
        self.partials.as_string()
    }

    /// Load from a file, using toml-rs and serialize
    pub fn load(path : &Path) -> Result<Self, Box<Error>> {
        let mut file : File = try!(File::open(path));
        let mut s = String::new();
        let _ = try!(file.read_to_string(&mut s));
        let mut parser = toml::Parser::new(s.as_ref());

        let parsed = parser.parse();
        let values = match parsed {
            Some(v) => v,
            None => {
                // We can unwrap here, becase parser.parse() == None means there were errors
                let err : Box<Error> = parser.errors.pop().unwrap().into();
                return Err(err)
            }
        };
        let table : toml::Value = toml::Value::Table(values);
        let mut decoder = toml::Decoder::new(table);
        let p = try!(rustc_serialize::Decodable::decode(&mut decoder));
        Ok(p)
    }
}

impl Default for Palette {
    fn default() -> Self {
        Palette {
            default_colors : FromIterator::from_iter(DEFAULT_COLORS.iter().map(
                |&(r,g,b)|{Color(r,g,b)}
            )),
            partials : PartialIDer::new(8, true),
            assigned : HashMap::new(),
            next_color : 0
        }
    }
}

#[derive(Serialize, RustcEncodable, Copy, Clone)]
struct AssignmentRef<'a> {
    names : &'a ObjectID,
    color : &'a Color
}

fn to_assignments<'a>(assigned : &'a HashMap<ObjectID, Color>) -> Vec<AssignmentRef<'a>> {
    assigned.into_iter().map(|(k, v)| {
        AssignmentRef{
            names : k,
            color : v
        }
    }).collect()
}

/// A reference to a palette, with minimal data, for serialization purposes
#[derive(Serialize, RustcEncodable, Clone)]
pub struct PaletteRef<'a> {
    defaults : &'a Vec<Color>,
    partials : &'a PartialIDer,
    assigned : Vec<AssignmentRef<'a>>
}

impl<'a> From<&'a Palette> for PaletteRef<'a> {
    fn from(palette: &'a Palette) -> Self {
        PaletteRef {
            defaults : &palette.default_colors,
            partials : &palette.partials,
            assigned : to_assignments(&palette.assigned),
        }
    }
}

#[derive(Deserialize, RustcDecodable, Clone)]
struct Assignment {
    names : ObjectID,
    color : Color
}

fn from_assignments(assigned : Vec<Assignment>) -> HashMap<ObjectID, Color> {
    assigned.into_iter().map(|a| {
        (a.names, a.color)
    }).collect()
}

/// A reference to a palette, with minimal data
#[derive(Deserialize, RustcDecodable, Clone)]
pub struct PaletteBasic {
    defaults : Vec<Color>,
    partials : PartialIDer,
    assigned : Option<Vec<Assignment>>,
    next_color : Option<usize>
}

impl From<PaletteBasic> for Palette {
    fn from(p : PaletteBasic) -> Self {
        Palette {
            default_colors : p.defaults,
            partials : p.partials,
            assigned : p.assigned.map(from_assignments).unwrap_or_default(),
            next_color : p.next_color.unwrap_or(0)
        }
    }
}

//TODO: add test for round-trip through toml and json

impl serde::Serialize for Color {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error> where S: serde::Serializer {
        let &Color(r,g,b) = self;
        serde::Serialize::serialize(&(r,g,b), serializer)
    }
}

impl serde::Deserialize for Color {
    fn deserialize<D : serde::Deserializer>(deserializer: &mut D) -> Result<Self, D::Error> {
        let (r,g,b) = try!(serde::Deserialize::deserialize(deserializer));
        Ok(Color(r,g,b))
    }
}

impl rustc_serialize::Encodable for Color {
    fn encode<S: rustc_serialize::Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        let &Color(r,g,b) = self;
        rustc_serialize::Encodable::encode(&(r,g,b), s)
    }
}

impl rustc_serialize::Decodable for Color {
    fn decode<D: rustc_serialize::Decoder>(d: &mut D) -> Result<Self, D::Error> {
        let (r,g,b) = try!(rustc_serialize::Decodable::decode(d));
        Ok(Color(r,g,b))
    }
}

impl serde::Serialize for Palette {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error> where S: serde::Serializer {
        PaletteRef::from(self).serialize(serializer)
    }
}

impl serde::Deserialize for Palette {
    fn deserialize<D : serde::Deserializer>(deserializer: &mut D) -> Result<Self, D::Error>  {
        let p = try!(PaletteBasic::deserialize(deserializer));
        Ok(Palette::from(p))
    }
}

impl rustc_serialize::Encodable for Palette {
    fn encode<S: rustc_serialize::Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        PaletteRef::from(self).encode(s)
    }
}

impl rustc_serialize::Decodable for Palette {
    fn decode<D: rustc_serialize::Decoder>(d: &mut D) -> Result<Self, D::Error> {
        let p = try!(PaletteBasic::decode(d));
        Ok(Palette::from(p))
    }
}
<|endoftext|>


<commit_before>use failure::{Error, ResultExt};
use mal::{self, MAL};
use regex::Regex;
use process;
use serde_json;
use std;
use std::collections::HashMap;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::process::ExitStatus;

#[derive(Fail, Debug)]
pub enum SeriesError {
    #[fail(display = "multiple series found")] MultipleSeriesFound,
    #[fail(display = "episode {} not found", _0)] EpisodeNotFound(u32),
    #[fail(display = "season {} information not found", _0)] SeasonInfoNotFound(u32),
}

#[derive(Debug)]
pub struct Series {
    pub name: String,
    pub data: SeriesData,
    pub episodes: HashMap<u32, PathBuf>,
    data_path: PathBuf,
}

impl Series {
    pub const DATA_FILE_NAME: &'static str = ".anitrack";

    pub fn from_path(path: &Path) -> Result<Series, Error> {
        let data_path = PathBuf::from(path).join(Series::DATA_FILE_NAME);
        let data = SeriesData::from_path_or_default(&data_path)?;

        let ep_data = EpisodeData::parse(path)?;

        Ok(Series {
            name: ep_data.series_name,
            data,
            episodes: ep_data.paths,
            data_path,
        })
    }

    pub fn play_episode(&self, ep_num: u32) -> Result<ExitStatus, Error> {
        let path = self.episodes.get(&ep_num)
            .ok_or(SeriesError::EpisodeNotFound(ep_num))?;

        let output = process::open_with_default(path).output()?;
        Ok(output.status)
    }

    pub fn set_season_data(&mut self, season: u32, info: SeasonInfo) {
        self.data.seasons.insert(season, info);
    }

    pub fn get_season_data(&self, season: u32) -> Result<&SeasonInfo, SeriesError> {
        self.data.seasons.get(&season).ok_or(SeriesError::SeasonInfoNotFound(season))
    }

    pub fn has_season_data(&self, season: u32) -> bool {
        self.data.seasons.contains_key(&season)
    }

    pub fn save_data(&self) -> Result<(), Error> {
        self.data.write_to(&self.data_path)
    }
}

#[derive(Debug, Default, Serialize, Deserialize)]
pub struct SeriesData {
    pub seasons: HashMap<u32, SeasonInfo>,
}

impl SeriesData {
    fn from_path(path: &Path) -> Result<SeriesData, Error> {
        let file = File::open(path)?;
        let data = serde_json::from_reader(file)?;

        Ok(data)
    }

    fn from_path_or_default(path: &Path) -> Result<SeriesData, Error> {
        if path.exists() {
            SeriesData::from_path(path)
        } else {
            Ok(SeriesData::default())
        }
    }

    fn write_to(&self, path: &Path) -> Result<(), Error> {
        let file = File::create(path)?;
        serde_json::to_writer_pretty(file, self)?;

        Ok(())
    }
}

#[derive(Fail, Debug)]
pub enum SeasonInfoError {
    #[fail(display = "no anime with id {} found with name [{}] on MAL", _0, _1)]
    UnknownAnimeID(u32, String),
}

#[derive(Debug, Serialize, Deserialize)]
pub struct SeasonInfo {
    pub series_id: u32,
    pub episodes: u32,
    pub search_title: String,
}

impl SeasonInfo {
    pub fn new(id: u32, episodes: u32, search_title: String) -> SeasonInfo {
        SeasonInfo {
            series_id: id,
            episodes,
            search_title,
        }
    }

    pub fn request_mal_info(&self, mal: &MAL) -> Result<mal::SeriesInfo, Error> {
        mal.search(&self.search_title)
            .context("MAL search failed")?
            .into_iter()
            .find(|i| i.id == self.series_id)
            .ok_or(SeasonInfoError::UnknownAnimeID(self.series_id, self.search_title.clone()).into())
    }
}

#[derive(Fail, Debug)]
pub enum EpisodeDataError {
    #[fail(display = "no episodes found")]
    NoEpisodesFound,
}

#[derive(Debug)]
pub struct EpisodeData {
    pub series_name: String,
    pub paths: HashMap<u32, PathBuf>,
}

impl EpisodeData {
    pub fn parse(path: &Path) -> Result<EpisodeData, Error> {
        let mut series = None;
        let mut episodes = HashMap::new();

        for entry in std::fs::read_dir(path)? {
            let path = entry?.path();

            let info = match EpisodeInfo::parse(&path) {
                Some(info) => info,
                None => continue,
            };

            if let Some(ref series) = series {
                if series != &info.series {
                    bail!(SeriesError::MultipleSeriesFound);
                }
            } else {
                series = Some(info.series);
            }

            episodes.insert(info.episode, path);
        }

        let series = series.ok_or(EpisodeDataError::NoEpisodesFound)?;

        Ok(EpisodeData {
            series_name: series,
            paths: episodes,
        })
    }
}

#[derive(Debug)]
struct EpisodeInfo {
    series: String,
    episode: u32,
}

impl EpisodeInfo {
    fn parse(path: &Path) -> Option<EpisodeInfo> {
        if !path.is_file() {
            return None;
        }

        lazy_static! {
            static ref EP_FORMAT: Regex = Regex::new(r"(?:\[.+?\]\s*)?(?P<series>.+?)\s*-?\s*(?P<episode>\d+).*?\..+?")
                .unwrap();
        }

        // Replace certain special characters with spaces since they can either
        // affect parsing or prevent finding results on MAL
        let filename = path.file_name()?.to_str()
            .unwrap()
            .replace('_', " ");

        let caps = EP_FORMAT.captures(&filename)?;

        let clean_name = {
            let raw = &caps["series"];
            raw.replace('.', " ")
               .replace(" -", ":") // Dashes typically represent a colon in file names
               .trim()
               .to_string()
        };

        let info = EpisodeInfo {
            series: clean_name,
            episode: caps["episode"].parse().ok()?,
        };

        Some(info)
    }
}<commit_msg>Optimize episode parsing regex to save a whopping 2 steps<commit_after>use failure::{Error, ResultExt};
use mal::{self, MAL};
use regex::Regex;
use process;
use serde_json;
use std;
use std::collections::HashMap;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::process::ExitStatus;

#[derive(Fail, Debug)]
pub enum SeriesError {
    #[fail(display = "multiple series found")] MultipleSeriesFound,
    #[fail(display = "episode {} not found", _0)] EpisodeNotFound(u32),
    #[fail(display = "season {} information not found", _0)] SeasonInfoNotFound(u32),
}

#[derive(Debug)]
pub struct Series {
    pub name: String,
    pub data: SeriesData,
    pub episodes: HashMap<u32, PathBuf>,
    data_path: PathBuf,
}

impl Series {
    pub const DATA_FILE_NAME: &'static str = ".anitrack";

    pub fn from_path(path: &Path) -> Result<Series, Error> {
        let data_path = PathBuf::from(path).join(Series::DATA_FILE_NAME);
        let data = SeriesData::from_path_or_default(&data_path)?;

        let ep_data = EpisodeData::parse(path)?;

        Ok(Series {
            name: ep_data.series_name,
            data,
            episodes: ep_data.paths,
            data_path,
        })
    }

    pub fn play_episode(&self, ep_num: u32) -> Result<ExitStatus, Error> {
        let path = self.episodes.get(&ep_num)
            .ok_or(SeriesError::EpisodeNotFound(ep_num))?;

        let output = process::open_with_default(path).output()?;
        Ok(output.status)
    }

    pub fn set_season_data(&mut self, season: u32, info: SeasonInfo) {
        self.data.seasons.insert(season, info);
    }

    pub fn get_season_data(&self, season: u32) -> Result<&SeasonInfo, SeriesError> {
        self.data.seasons.get(&season).ok_or(SeriesError::SeasonInfoNotFound(season))
    }

    pub fn has_season_data(&self, season: u32) -> bool {
        self.data.seasons.contains_key(&season)
    }

    pub fn save_data(&self) -> Result<(), Error> {
        self.data.write_to(&self.data_path)
    }
}

#[derive(Debug, Default, Serialize, Deserialize)]
pub struct SeriesData {
    pub seasons: HashMap<u32, SeasonInfo>,
}

impl SeriesData {
    fn from_path(path: &Path) -> Result<SeriesData, Error> {
        let file = File::open(path)?;
        let data = serde_json::from_reader(file)?;

        Ok(data)
    }

    fn from_path_or_default(path: &Path) -> Result<SeriesData, Error> {
        if path.exists() {
            SeriesData::from_path(path)
        } else {
            Ok(SeriesData::default())
        }
    }

    fn write_to(&self, path: &Path) -> Result<(), Error> {
        let file = File::create(path)?;
        serde_json::to_writer_pretty(file, self)?;

        Ok(())
    }
}

#[derive(Fail, Debug)]
pub enum SeasonInfoError {
    #[fail(display = "no anime with id {} found with name [{}] on MAL", _0, _1)]
    UnknownAnimeID(u32, String),
}

#[derive(Debug, Serialize, Deserialize)]
pub struct SeasonInfo {
    pub series_id: u32,
    pub episodes: u32,
    pub search_title: String,
}

impl SeasonInfo {
    pub fn new(id: u32, episodes: u32, search_title: String) -> SeasonInfo {
        SeasonInfo {
            series_id: id,
            episodes,
            search_title,
        }
    }

    pub fn request_mal_info(&self, mal: &MAL) -> Result<mal::SeriesInfo, Error> {
        mal.search(&self.search_title)
            .context("MAL search failed")?
            .into_iter()
            .find(|i| i.id == self.series_id)
            .ok_or(SeasonInfoError::UnknownAnimeID(self.series_id, self.search_title.clone()).into())
    }
}

#[derive(Fail, Debug)]
pub enum EpisodeDataError {
    #[fail(display = "no episodes found")]
    NoEpisodesFound,
}

#[derive(Debug)]
pub struct EpisodeData {
    pub series_name: String,
    pub paths: HashMap<u32, PathBuf>,
}

impl EpisodeData {
    pub fn parse(path: &Path) -> Result<EpisodeData, Error> {
        let mut series = None;
        let mut episodes = HashMap::new();

        for entry in std::fs::read_dir(path)? {
            let path = entry?.path();

            let info = match EpisodeInfo::parse(&path) {
                Some(info) => info,
                None => continue,
            };

            if let Some(ref series) = series {
                if series != &info.series {
                    bail!(SeriesError::MultipleSeriesFound);
                }
            } else {
                series = Some(info.series);
            }

            episodes.insert(info.episode, path);
        }

        let series = series.ok_or(EpisodeDataError::NoEpisodesFound)?;

        Ok(EpisodeData {
            series_name: series,
            paths: episodes,
        })
    }
}

#[derive(Debug)]
struct EpisodeInfo {
    series: String,
    episode: u32,
}

impl EpisodeInfo {
    fn parse(path: &Path) -> Option<EpisodeInfo> {
        if !path.is_file() {
            return None;
        }

        lazy_static! {
            static ref EP_FORMAT: Regex = Regex::new(r"(?:\[.+?\]\s*)?(?P<series>.+?)\s*(?:-\s*)?(?P<episode>\d+).*?\..+?")
                .unwrap();
        }

        // Replace certain special characters with spaces since they can either
        // affect parsing or prevent finding results on MAL
        let filename = path.file_name()?.to_str()
            .unwrap()
            .replace('_', " ");

        let caps = EP_FORMAT.captures(&filename)?;

        let clean_name = {
            let raw = &caps["series"];
            raw.replace('.', " ")
               .replace(" -", ":") // Dashes typically represent a colon in file names
               .trim()
               .to_string()
        };

        let info = EpisodeInfo {
            series: clean_name,
            episode: caps["episode"].parse().ok()?,
        };

        Some(info)
    }
}<|endoftext|>