<commit_before>use std::borrow::Borrow;
use std::iter::FromIterator;
use std::slice::{Iter, IterMut};
use std::vec::IntoIter;

use crate::stable_hasher::{HashStable, StableHasher};

/// A map type implemented as a vector of pairs `K` (key) and `V` (value).
/// It currently provides a subset of all the map operations, the rest could be added as needed.
#[derive(Clone, Encodable, Decodable, Debug)]
pub struct VecMap<K, V>(Vec<(K, V)>);

impl<K, V> VecMap<K, V>
where
    K: PartialEq,
{
    pub fn new() -> Self {
        VecMap(Default::default())
    }

    /// Sets the value of the entry, and returns the entry's old value.
    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
        if let Some(elem) = self.0.iter_mut().find(|(key, _)| *key == k) {
            Some(std::mem::replace(&mut elem.1, v))
        } else {
            self.0.push((k, v));
            None
        }
    }

    /// Gets a reference to the value in the entry.
    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
    where
        K: Borrow<Q>,
        Q: Eq,
    {
        self.0.iter().find(|(key, _)| k == key.borrow()).map(|elem| &elem.1)
    }

    /// Returns the value corresponding to the supplied predicate filter.
    ///
    /// The supplied predicate will be applied to each (key, value) pair and it will return a
    /// reference to the values where the predicate returns `true`.
    pub fn get_by(&self, mut predicate: impl FnMut(&(K, V)) -> bool) -> Option<&V> {
        self.0.iter().find(|kv| predicate(kv)).map(|elem| &elem.1)
    }

    /// Returns `true` if the map contains a value for the specified key.
    ///
    /// The key may be any borrowed form of the map's key type,
    /// [`Eq`] on the borrowed form *must* match those for
    /// the key type.
    pub fn contains_key<Q: ?Sized>(&self, k: &Q) -> bool
    where
        K: Borrow<Q>,
        Q: Eq,
    {
        self.get(k).is_some()
    }

    /// Returns `true` if the map contains no elements.
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    pub fn iter(&self) -> Iter<'_, (K, V)> {
        self.into_iter()
    }

    pub fn iter_mut(&mut self) -> IterMut<'_, (K, V)> {
        self.into_iter()
    }
}

impl<K, V> Default for VecMap<K, V> {
    #[inline]
    fn default() -> Self {
        Self(Default::default())
    }
}

impl<K, V> From<Vec<(K, V)>> for VecMap<K, V> {
    fn from(vec: Vec<(K, V)>) -> Self {
        Self(vec)
    }
}

impl<K, V> Into<Vec<(K, V)>> for VecMap<K, V> {
    fn into(self) -> Vec<(K, V)> {
        self.0
    }
}

impl<K, V> FromIterator<(K, V)> for VecMap<K, V> {
    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
        Self(iter.into_iter().collect())
    }
}

impl<'a, K, V> IntoIterator for &'a VecMap<K, V> {
    type Item = &'a (K, V);
    type IntoIter = Iter<'a, (K, V)>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.0.iter()
    }
}

impl<'a, K, V> IntoIterator for &'a mut VecMap<K, V> {
    type Item = &'a mut (K, V);
    type IntoIter = IterMut<'a, (K, V)>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.0.iter_mut()
    }
}

impl<K, V> IntoIterator for VecMap<K, V> {
    type Item = (K, V);
    type IntoIter = IntoIter<(K, V)>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

impl<K, V> Extend<(K, V)> for VecMap<K, V> {
    fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
        self.0.extend(iter);
    }

    fn extend_one(&mut self, item: (K, V)) {
        self.0.extend_one(item);
    }

    fn extend_reserve(&mut self, additional: usize) {
        self.0.extend_reserve(additional);
    }
}

impl<K, V, CTX> HashStable<CTX> for VecMap<K, V>
where
    K: HashStable<CTX> + Eq,
    V: HashStable<CTX>,
{
    fn hash_stable(&self, hcx: &mut CTX, hasher: &mut StableHasher) {
        self.0.hash_stable(hcx, hasher)
    }
}

#[cfg(test)]
mod tests;
<commit_msg>Fix VecMap Extend impl<commit_after>use std::borrow::Borrow;
use std::iter::FromIterator;
use std::slice::{Iter, IterMut};
use std::vec::IntoIter;

use crate::stable_hasher::{HashStable, StableHasher};

/// A map type implemented as a vector of pairs `K` (key) and `V` (value).
/// It currently provides a subset of all the map operations, the rest could be added as needed.
#[derive(Clone, Encodable, Decodable, Debug)]
pub struct VecMap<K, V>(Vec<(K, V)>);

impl<K, V> VecMap<K, V>
where
    K: PartialEq,
{
    pub fn new() -> Self {
        VecMap(Default::default())
    }

    /// Sets the value of the entry, and returns the entry's old value.
    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
        if let Some(elem) = self.0.iter_mut().find(|(key, _)| *key == k) {
            Some(std::mem::replace(&mut elem.1, v))
        } else {
            self.0.push((k, v));
            None
        }
    }

    /// Gets a reference to the value in the entry.
    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
    where
        K: Borrow<Q>,
        Q: Eq,
    {
        self.0.iter().find(|(key, _)| k == key.borrow()).map(|elem| &elem.1)
    }

    /// Returns the value corresponding to the supplied predicate filter.
    ///
    /// The supplied predicate will be applied to each (key, value) pair and it will return a
    /// reference to the values where the predicate returns `true`.
    pub fn get_by(&self, mut predicate: impl FnMut(&(K, V)) -> bool) -> Option<&V> {
        self.0.iter().find(|kv| predicate(kv)).map(|elem| &elem.1)
    }

    /// Returns `true` if the map contains a value for the specified key.
    ///
    /// The key may be any borrowed form of the map's key type,
    /// [`Eq`] on the borrowed form *must* match those for
    /// the key type.
    pub fn contains_key<Q: ?Sized>(&self, k: &Q) -> bool
    where
        K: Borrow<Q>,
        Q: Eq,
    {
        self.get(k).is_some()
    }

    /// Returns `true` if the map contains no elements.
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    pub fn iter(&self) -> Iter<'_, (K, V)> {
        self.into_iter()
    }

    pub fn iter_mut(&mut self) -> IterMut<'_, (K, V)> {
        self.into_iter()
    }
}

impl<K, V> Default for VecMap<K, V> {
    #[inline]
    fn default() -> Self {
        Self(Default::default())
    }
}

impl<K, V> From<Vec<(K, V)>> for VecMap<K, V> {
    fn from(vec: Vec<(K, V)>) -> Self {
        Self(vec)
    }
}

impl<K, V> Into<Vec<(K, V)>> for VecMap<K, V> {
    fn into(self) -> Vec<(K, V)> {
        self.0
    }
}

impl<K, V> FromIterator<(K, V)> for VecMap<K, V> {
    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
        Self(iter.into_iter().collect())
    }
}

impl<'a, K, V> IntoIterator for &'a VecMap<K, V> {
    type Item = &'a (K, V);
    type IntoIter = Iter<'a, (K, V)>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.0.iter()
    }
}

impl<'a, K, V> IntoIterator for &'a mut VecMap<K, V> {
    type Item = &'a mut (K, V);
    type IntoIter = IterMut<'a, (K, V)>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.0.iter_mut()
    }
}

impl<K, V> IntoIterator for VecMap<K, V> {
    type Item = (K, V);
    type IntoIter = IntoIter<(K, V)>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

impl<K: PartialEq, V> Extend<(K, V)> for VecMap<K, V> {
    fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
        for (k, v) in iter {
            self.insert(k, v);
        }
    }

    fn extend_one(&mut self, (k, v): (K, V)) {
        self.insert(k, v);
    }

    fn extend_reserve(&mut self, additional: usize) {
        self.0.extend_reserve(additional);
    }
}

impl<K, V, CTX> HashStable<CTX> for VecMap<K, V>
where
    K: HashStable<CTX> + Eq,
    V: HashStable<CTX>,
{
    fn hash_stable(&self, hcx: &mut CTX, hasher: &mut StableHasher) {
        self.0.hash_stable(hcx, hasher)
    }
}

#[cfg(test)]
mod tests;
<|endoftext|>


<commit_before>#[cfg(feature = "bindgen")]
extern crate bindgen;
extern crate pkg_config;

use pkg_config::{Config, Library};
use std::env;
use std::path::PathBuf;

fn find_gdal() -> Option<Library> {
    Config::new().probe("gdal").ok()
}

#[cfg(feature = "bindgen")]
pub fn build(gdal: Option<Library>, out_path: &std::path::Path) {
    #[cfg(windows)]
    {
        // get the path to GDAL_HOME
        let home_path = env::var("GDAL_HOME").expect("Environment variable $GDAL_HOME not found!");

        // detect the path to gdal_i.lib (works for MSVC and GNU)
        let lib_suffix = "_i";
        let lib_search_path = Path::new(&home_path).join("lib");
        let lib_path = lib_search_path.join(&format!("{}{}.lib", lib_name, lib_suffix));

        if lib_search_path.exists() && lib_path.exists() {
            println!(
                "cargo:rustc-link-search={}",
                lib_search_path.to_string_lossy()
            );
            println!(
                "cargo:rustc-link-lib={}={}",
                link_type,
                format!("{}{}", lib_name, lib_suffix)
            );
        } else {
            #[cfg(target_env = "msvc")]
            {
                panic!("windows-msvc requires gdal_i.lib to be found in $GDAL_HOME\\lib.");
            }

            #[cfg(target_env = "gnu")]
            {
                // detect if a gdal{version}.dll is available
                let versions = [201, 200, 111, 110];
                let bin_path = Path::new(&home_path).join("bin");
                if let Some(version) = versions
                    .iter()
                    .find(|v| bin_path.join(&format!("{}{}.dll", lib_name, v)).exists())
                {
                    println!("cargo:rustc-link-search={}", bin_path.to_string_lossy());
                    println!(
                        "cargo:rustc-link-lib={}={}",
                        link_type,
                        format!("{}{}", lib_name, version)
                    );
                } else {
                    panic!("windows-gnu requires either gdal_i.lib in $GDAL_HOME\\lib OR gdal{version}.dll in $GDAL_HOME\\bin.");
                }
            }
        }
    }

    let mut builder = bindgen::Builder::default();
    if let Some(gdal) = gdal {
        for path in &gdal.include_paths {
            builder = builder.clang_arg("-I");
            builder = builder.clang_arg(path.to_str().unwrap());
        }
    }

    builder
        .header("wrapper.h")
        .constified_enum_module(".*")
        .ctypes_prefix("libc")
        .whitelist_function("CPL.*")
        .whitelist_function("GDAL.*")
        .whitelist_function("OGR.*")
        .whitelist_function("OSR.*")
        .whitelist_function("OCT.*")
        .whitelist_function("VSI.*")
        .generate()
        .expect("Unable to generate bindings")
        .write_to_file(out_path)
        .expect("Unable to write bindings to file");
}

fn main() {
    let out_path = PathBuf::from(env::var("OUT_DIR").unwrap()).join("bindings.rs");

    #[cfg(feature = "bindgen")]
    build(find_gdal(), &out_path);

    #[cfg(not(feature = "bindgen"))]
    {
        find_gdal();

        let prebuilt_paths = &[
            #[cfg(feature = "min_gdal_version_1_11")]
            "prebuilt-bindings/gdal_1.11.rs",
            #[cfg(feature = "min_gdal_version_2_0")]
            "prebuilt-bindings/gdal_2.0.rs",
            #[cfg(feature = "min_gdal_version_2_1")]
            "prebuilt-bindings/gdal_2.1.rs",
            #[cfg(feature = "min_gdal_version_2_2")]
            "prebuilt-bindings/gdal_2.2.rs",
        ];
        std::fs::copy(&prebuilt_paths[prebuilt_paths.len() - 1], &out_path)
            .expect("Can't copy bindings to output directory");
    }
}
<commit_msg>Try to fix Windows builds<commit_after>#[cfg(feature = "bindgen")]
extern crate bindgen;
extern crate pkg_config;

use pkg_config::Config;
use std::env;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};

#[cfg(feature = "bindgen")]
pub fn write_bindings(include_paths: Vec<String>, out_path: &Path) {
    let mut builder = bindgen::Builder::default()
        .header("wrapper.h")
        .constified_enum_module(".*")
        .ctypes_prefix("libc")
        .whitelist_function("CPL.*")
        .whitelist_function("GDAL.*")
        .whitelist_function("OGR.*")
        .whitelist_function("OSR.*")
        .whitelist_function("OCT.*")
        .whitelist_function("VSI.*");

    for path in include_paths {
        builder = builder.clang_arg("-I");
        builder = builder.clang_arg(path);
    }

    builder
        .generate()
        .expect("Unable to generate bindings")
        .write_to_file(out_path)
        .expect("Unable to write bindings to file");
}

fn env_dir(var: &str) -> Option<PathBuf> {
    let dir = env::var_os(var).map(PathBuf::from);

    if let Some(ref dir) = dir {
        if !dir.exists() {
            panic!("{} was set to {}, which doesn't exist.", var, dir.display());
        }
    }

    dir
}

fn find_gdal_dll(lib_dir: &Path) -> io::Result<Option<String>> {
    for e in fs::read_dir(&lib_dir)? {
        let e = e?;
        let name = e.file_name();
        let name = name.to_str().unwrap();
        if name.starts_with("gdal") && name.ends_with(".dll") {
            return Ok(Some(String::from(name)));
        }
    }
    Ok(None)
}

fn main() {
    println!("cargo:rerun-if-env-changed=GDAL_STATIC");
    println!("cargo:rerun-if-env-changed=GDAL_INCLUDE_DIR");
    println!("cargo:rerun-if-env-changed=GDAL_LIB_DIR");
    println!("cargo:rerun-if-env-changed=GDAL_HOME");

    let mut need_metadata = true;
    let mut lib_name = String::from("gdal");

    let mut prefer_static =
        env::var_os("GDAL_STATIC").is_some() && env::var_os("GDAL_DYNAMIC").is_none();

    let mut include_dir = env_dir("GDAL_INCLUDE_DIR");
    let mut lib_dir = env_dir("GDAL_LIB_DIR");
    let home_dir = env_dir("GDAL_HOME");

    let mut found = false;
    if cfg!(windows) {
        // first, look for a static library in $GDAL_LIB_DIR or $GDAL_HOME/lib
        // works in windows-msvc and windows-gnu
        if let Some(ref lib_dir) = lib_dir {
            let lib_path = lib_dir.join("gdal_i.lib");
            if lib_path.exists() {
                prefer_static = true;
                lib_name = String::from("gdal_i");
                found = true;
            }
        }
        if !found {
            if let Some(ref home_dir) = home_dir {
                let home_lib_dir = home_dir.join("lib");
                let lib_path = home_lib_dir.join("gdal_i.lib");
                if lib_path.exists() {
                    prefer_static = true;
                    lib_name = String::from("gdal_i");
                    lib_dir = Some(home_lib_dir);
                    found = true;
                }
            }
        }
        if !found {
            if cfg!(target_env = "msvc") {
                panic!("windows-gnu requires gdal_i.lib to be present in either $GDAL_LIB_DIR or $GDAL_HOME\\lib.");
            }

            // otherwise, look for a gdalxxx.dll in $GDAL_HOME/bin
            // works in windows-gnu
            if let Some(ref home_dir) = home_dir {
                let bin_dir = home_dir.join("bin");
                if bin_dir.exists() {
                    if let Some(name) = find_gdal_dll(&bin_dir).unwrap() {
                        prefer_static = false;
                        lib_dir = Some(bin_dir);
                        lib_name = name;
                    }
                }
            }
        }
    }

    if let Some(ref home_dir) = home_dir {
        if include_dir.is_none() {
            let dir = home_dir.join("include");
            if cfg!(feature = "bindgen") && !dir.exists() {
                panic!(
                    "bindgen was enabled, but GDAL_INCLUDE_DIR was not set and {} doesn't exist.",
                    dir.display()
                );
            }
            include_dir = Some(dir);
        }

        if lib_dir.is_none() {
            let dir = home_dir.join("lib");
            if !dir.exists() {
                panic!(
                    "GDAL_LIB_DIR was not set and {} doesn't exist.",
                    dir.display()
                );
            }
            lib_dir = Some(dir);
        }
    }

    if let Some(lib_dir) = lib_dir {
        let link_type = if prefer_static { "static" } else { "dylib" };

        println!("cargo:rustc-link-lib={}={}", link_type, lib_name);
        println!("cargo:rustc-link-search={}", lib_dir.to_str().unwrap());

        need_metadata = false;
    }

    let out_path = PathBuf::from(env::var("OUT_DIR").unwrap()).join("bindings.rs");

    let mut include_paths = Vec::new();
    if let Some(ref dir) = include_dir {
        include_paths.push(dir.as_path().to_str().unwrap().to_string());
    }

    let gdal = Config::new()
        .statik(prefer_static)
        .cargo_metadata(need_metadata)
        .probe("gdal");

    if let Ok(gdal) = gdal {
        for dir in gdal.include_paths {
            include_paths.push(dir.to_str().unwrap().to_string());
        }
    }

    #[cfg(feature = "bindgen")]
    write_bindings(include_paths, &out_path);

    #[cfg(not(feature = "bindgen"))]
    {
        let prebuilt_paths = &[
            #[cfg(feature = "min_gdal_version_1_11")]
            "prebuilt-bindings/gdal_1.11.rs",
            #[cfg(feature = "min_gdal_version_2_0")]
            "prebuilt-bindings/gdal_2.0.rs",
            #[cfg(feature = "min_gdal_version_2_1")]
            "prebuilt-bindings/gdal_2.1.rs",
            #[cfg(feature = "min_gdal_version_2_2")]
            "prebuilt-bindings/gdal_2.2.rs",
        ];
        std::fs::copy(&prebuilt_paths[prebuilt_paths.len() - 1], &out_path)
            .expect("Can't copy bindings to output directory");
    }
}
<|endoftext|>


<commit_before>
use rand::{thread_rng, Rng};

use std::io;
use std::fmt;

use opcode::{Opcode, OpcodeError, SetRegMode};

const PROGRAM_START: u16 = 0x200;
const FONT_START: u16 = 0x50;

const MEMORY_SIZE: usize = 4096;
const REGISTER_COUNT: usize = 16;

// Thanks to: http://www.multigesture.net/articles/how-to-write-an-emulator-chip-8-interpreter/
const FONTMAP: [u8; 80] = [
  0xF0, 0x90, 0x90, 0x90, 0xF0, // 0
  0x20, 0x60, 0x20, 0x20, 0x70, // 1
  0xF0, 0x10, 0xF0, 0x80, 0xF0, // 2
  0xF0, 0x10, 0xF0, 0x10, 0xF0, // 3
  0x90, 0x90, 0xF0, 0x10, 0x10, // 4
  0xF0, 0x80, 0xF0, 0x10, 0xF0, // 5
  0xF0, 0x80, 0xF0, 0x90, 0xF0, // 6
  0xF0, 0x10, 0x20, 0x40, 0x40, // 7
  0xF0, 0x90, 0xF0, 0x90, 0xF0, // 8
  0xF0, 0x90, 0xF0, 0x10, 0xF0, // 9
  0xF0, 0x90, 0xF0, 0x90, 0x90, // A
  0xE0, 0x90, 0xE0, 0x90, 0xE0, // B
  0xF0, 0x80, 0x80, 0x80, 0xF0, // C
  0xE0, 0x90, 0x90, 0x90, 0xE0, // D
  0xF0, 0x80, 0xF0, 0x80, 0xF0, // E
  0xF0, 0x80, 0xF0, 0x80, 0x80  // F
];

#[derive(Debug)]
pub enum RuntimeError {
    EmptyCallStack,
    InvalidRegister(u8),
    AddressOutOfBounds(u16),
    OpcodeErr(OpcodeError),
}

pub struct Chip8 {
    pub memory: [u8; MEMORY_SIZE],
    pub regs: [u8; REGISTER_COUNT], // registers V0 - V15
    pub addressReg: u16, // register I

    pub pc: u16,
    pub stack: Vec<u16>,

    pub delay_timer: u16,
    pub sound_timer: u16,

    pub screen: [[bool; 64]; 32],
}

impl Clone for Chip8 {
    fn clone(&self) -> Chip8 {
        let mut memory = [0; 4096];
        for (offset, byte) in self.memory.iter().enumerate() {
            memory[offset] = *byte;
        }

        let mut regs = [0; 16];
        for (offset, reg) in self.regs.iter().enumerate() {
            regs[offset] = *reg;
        }

        let mut screen = [[false; 64]; 32];
        for (y, row) in self.screen.iter().enumerate() {
            for (x, elem) in row.iter().enumerate() {
                screen[y][x] = *elem;
            }
        }

        Chip8 {
            memory: memory,
            regs: regs,
            addressReg: self.addressReg,

            pc: self.pc,
            stack: self.stack.clone(),

            delay_timer: self.delay_timer,
            sound_timer: self.sound_timer,

            screen: screen,
        }
    }
}

impl Chip8 {
    pub fn new() -> Chip8 {
        let mut chip8 = Chip8 {
            memory: [0; 4096],
            regs: [0; 16],
            addressReg: 0,

            pc: PROGRAM_START,
            stack: Vec::new(),

            delay_timer: 0,
            sound_timer: 0,

            screen: [[false; 64]; 32],
        };

        chip8.inject_fontmap();
        chip8
    }

    pub fn inject_fontmap(&mut self) {
        for (offset, byte) in FONTMAP.iter().enumerate() {
            self.memory[FONT_START as usize + offset] = *byte;
        }
    }

    pub fn load_program<R: io::Read>(&mut self, mut program: R) -> io::Result<()> {
        let mut bytes = Vec::new();
        try!(program.read_to_end(&mut bytes));

        for (offset, byte) in bytes.iter().enumerate() {
            self.memory[offset + PROGRAM_START as usize] = *byte;
        }

        Ok(())
    }

    pub fn cycle(&mut self, keys: [bool; 16]) -> Result<(), RuntimeError> {
        use self::RuntimeError::*;

        Chip8::decrement_timer(&mut self.delay_timer);
        Chip8::decrement_timer(&mut self.sound_timer);

        let pc_index = self.pc as usize;
        let opcode_bytes = (self.memory[pc_index] as u16) << 8 | (self.memory[pc_index + 1] as u16);

        let opcode = match Opcode::from_u16(opcode_bytes) {
            Ok(opcode) => opcode,
            Err(err) => return Err(OpcodeErr(err)),
        };
        println!("{:X}: {:?}", opcode_bytes, opcode);

        self.pc += 2;
        try!(self.execute_opcode(opcode, keys));

        Ok(())
    }

    pub fn clear_screen(&mut self) {
        for row in self.screen.iter_mut() {
            for col in row.iter_mut() {
                *col = false;
            }
        }
    }

    // Wrapping is performed in this function, no need to perform it outside
    // Returns if a pixel was unset
    pub fn set_pixel(&mut self, x: usize, y: usize) -> bool {
        // Equivalent to using the mod operator, but faster
        let x = x & 63;
        let y = y & 31;

        let previous_state = self.screen[y][x];
        self.screen[y][x] = !self.screen[y][x];

        // return whether a pixel was previously set and then now unset
        previous_state
    }

    pub fn execute_opcode(&mut self, opcode: Opcode, keys: [bool; 16]) -> Result<(), RuntimeError> {
        use self::RuntimeError::*;
        use opcode::Opcode::*;

        match opcode {
            ClearScreen => self.clear_screen(),
            Return => {
                self.pc = match self.stack.pop() {
                    Some(addr) => addr,
                    None => return Err(EmptyCallStack),
                };
            },

            JumpTo { addr, plus_v0 } => {
                self.pc = addr;

                if plus_v0 {
                    self.pc += self.regs[0] as u16;
                }
            },
            Call(addr) => {
                self.stack.push(self.pc);
                self.pc = addr;
            },

            SkipIfRegEqualConst { not_equal, reg, value } => {
                let mut should_jump = self.regs[reg as usize] == value;

                if not_equal {
                    should_jump = !should_jump; // Effectively computes self.regs[reg] != value
                }

                if should_jump {
                    self.pc += 2;
                }
            },
            SkipIfRegsEqual { not_equal, regs: (v_x, v_y) } => {
                let mut should_jump = self.regs[v_x as usize] == self.regs[v_y as usize];

                if not_equal {
                    should_jump = !should_jump;
                }

                if should_jump {
                    self.pc += 2;
                }
            },

            SetRegToConst { add, reg, value } => {
                if add {
                    let mut value = (self.regs[reg as usize] as u32 + value as u32) & 255;
                    self.regs[reg as usize] = value as u8;
                } else {
                    self.regs[reg as usize] = value;
                }
            },
            SetRegToReg { regs: (v_x, v_y), mode } => {
                let v_x = v_x as usize;
                let v_y = v_y as usize;

                match mode {
                    SetRegMode::Copy => self.regs[v_x] = self.regs[v_y],

                    SetRegMode::Or => self.regs[v_x] |= self.regs[v_y],
                    SetRegMode::And => self.regs[v_x] &= self.regs[v_y],
                    SetRegMode::Xor => self.regs[v_x] ^= self.regs[v_y],

                    SetRegMode::Add => {
                        self.regs[0xF] = 0;

                        let mut reg_value = self.regs[v_x] as usize + self.regs[v_y] as usize;
                        if reg_value > 255 {
                            reg_value -= 255;
                            self.regs[0xF] = 1;
                        }

                        self.regs[v_x] = reg_value as u8;
                    },
                    SetRegMode::Subtract | SetRegMode::InverseSubtract => {
                        self.regs[0xF] = 1;

                        let mut reg_value = if mode == SetRegMode::Subtract {
                            self.regs[v_x] as isize - self.regs[v_y] as isize
                        } else { // Must be InverseSubtract
                            self.regs[v_y] as isize - self.regs[v_x] as isize
                        };

                        if reg_value < 0 {
                            reg_value += 255;
                            self.regs[0xF] = 0;
                        }

                        self.regs[v_x] = reg_value as u8;
                    },

                    // v_y is ignored for the shift opcodes, not sure why
                    SetRegMode::ShiftLeft => {
                        self.regs[0xF] = self.regs[v_x] & 128;

                        self.regs[v_x] <<= 1;
                        println!("{}", self.regs[v_x]);
                    },
                    SetRegMode::ShiftRight => {
                        self.regs[0xF] = self.regs[v_x] & 0x1;
                        self.regs[v_x] >>= 1;
                    }
                }
            },

            SetAddressReg(addr) => self.addressReg = addr,
            SetRegToRandom { reg, mask } => {
                let rand: u8 = thread_rng().gen();
                self.regs[reg as usize] = rand & mask;
            },

            DrawSprite { regs: (v_x, v_y), rows } => {
                let x = self.regs[v_x as usize] as usize;
                let y = self.regs[v_y as usize] as usize;

                self.regs[0xF] = 0;

                for row in 0..rows {
                    let sprite_slice = self.memory[(self.addressReg + row as u16) as usize];

                    for col in 0..8 {
                        if (sprite_slice & (128 >> col)) != 0 {
                            if self.set_pixel(x + col as usize, y + row as usize) {
                                self.regs[0xF] = 1;
                            }
                        }
                    }
                }
            },

            SetRegToDelayTimer(reg) => self.regs[reg as usize] = self.delay_timer as u8,

            SetDelayTimerToReg(reg) => self.delay_timer = self.regs[reg as usize] as u16,
            SetSoundTimerToReg(reg) => self.sound_timer = self.regs[reg as usize] as u16,

            AddRegToAddressReg(reg) => self.addressReg += self.regs[reg as usize] as u16,
            SetAddressRegToCharInReg(reg) => {
                let ch = self.regs[reg as usize];
                self.addressReg = FONT_START + ch as u16 * 5;
            },

            WaitForKeyInReg(reg) => {
                // TODO
            },
            SkipIfKeyInRegPressed { not_pressed, reg } => {
                let mut should_jump = keys[self.regs[reg as usize] as usize];

                if not_pressed {
                    should_jump = !should_jump;
                }

                if should_jump {
                    self.pc += 2;
                }
            },

            // See http://en.wikipedia.org/wiki/Binary-coded_decimal
            // n mod 10 => Gets the ones digit out of a number
            RegToBCD(reg) => {
                let number = self.regs[reg as usize];

                let hundreds_digit = number / 100;
                let tens_digit = (number / 10) % 10; // Dividing by ten slides the tens digit into the ones digit
                let ones_digit = number % 10;

                self.memory[(self.addressReg) as usize] = hundreds_digit;
                self.memory[(self.addressReg + 1) as usize] = tens_digit;
                self.memory[(self.addressReg + 2) as usize] = ones_digit;
            },

            DumpRegsToAddr(reg) => {
                for cur_reg in 0..(reg + 1) {
                    self.memory[(self.addressReg + cur_reg as u16)  as usize] = self.regs[cur_reg as usize];
                }
            },
            LoadRegsFromAddr(reg) => {
                for cur_reg in 0..(reg + 1) {
                    self.regs[cur_reg as usize] = self.memory[(self.addressReg + cur_reg as u16) as usize];
                }
            }
        }

        Ok(())
    }

    fn decrement_timer(timer: &mut u16) {
        if *timer != 0 {
            *timer -= 1;
        }
    }
}

impl fmt::Debug for Chip8 {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        try!(writeln!(fmt, "Program Counter: 0x{:X} ({})", self.pc, self.pc));
        try!(writeln!(fmt, "Address Register: 0x{:X} ({})", self.addressReg, self.addressReg));
        try!(writeln!(fmt, "Stack: {:?}", self.stack));
        try!(writeln!(fmt, "Delay Timer: {}", self.delay_timer));
        try!(writeln!(fmt, "Sound Timer: {}", self.sound_timer));

        write!(fmt, "Register Contents: {:?}", self.regs)
    }
}
<commit_msg>Arthimetic now fully works, emulator is basically done ;)<commit_after>
use rand::{thread_rng, Rng};

use std::io;
use std::fmt;

use opcode::{Opcode, OpcodeError, SetRegMode};

const PROGRAM_START: u16 = 0x200;
const FONT_START: u16 = 0x50;

const MEMORY_SIZE: usize = 4096;
const REGISTER_COUNT: usize = 16;

// Thanks to: http://www.multigesture.net/articles/how-to-write-an-emulator-chip-8-interpreter/
const FONTMAP: [u8; 80] = [
  0xF0, 0x90, 0x90, 0x90, 0xF0, // 0
  0x20, 0x60, 0x20, 0x20, 0x70, // 1
  0xF0, 0x10, 0xF0, 0x80, 0xF0, // 2
  0xF0, 0x10, 0xF0, 0x10, 0xF0, // 3
  0x90, 0x90, 0xF0, 0x10, 0x10, // 4
  0xF0, 0x80, 0xF0, 0x10, 0xF0, // 5
  0xF0, 0x80, 0xF0, 0x90, 0xF0, // 6
  0xF0, 0x10, 0x20, 0x40, 0x40, // 7
  0xF0, 0x90, 0xF0, 0x90, 0xF0, // 8
  0xF0, 0x90, 0xF0, 0x10, 0xF0, // 9
  0xF0, 0x90, 0xF0, 0x90, 0x90, // A
  0xE0, 0x90, 0xE0, 0x90, 0xE0, // B
  0xF0, 0x80, 0x80, 0x80, 0xF0, // C
  0xE0, 0x90, 0x90, 0x90, 0xE0, // D
  0xF0, 0x80, 0xF0, 0x80, 0xF0, // E
  0xF0, 0x80, 0xF0, 0x80, 0x80  // F
];

#[derive(Debug)]
pub enum RuntimeError {
    EmptyCallStack,
    InvalidRegister(u8),
    AddressOutOfBounds(u16),
    OpcodeErr(OpcodeError),
}

pub struct Chip8 {
    pub memory: [u8; MEMORY_SIZE],
    pub regs: [u8; REGISTER_COUNT], // registers V0 - V15
    pub addressReg: u16, // register I

    pub pc: u16,
    pub stack: Vec<u16>,

    pub delay_timer: u16,
    pub sound_timer: u16,

    pub screen: [[bool; 64]; 32],
}

impl Clone for Chip8 {
    fn clone(&self) -> Chip8 {
        let mut memory = [0; 4096];
        for (offset, byte) in self.memory.iter().enumerate() {
            memory[offset] = *byte;
        }

        let mut regs = [0; 16];
        for (offset, reg) in self.regs.iter().enumerate() {
            regs[offset] = *reg;
        }

        let mut screen = [[false; 64]; 32];
        for (y, row) in self.screen.iter().enumerate() {
            for (x, elem) in row.iter().enumerate() {
                screen[y][x] = *elem;
            }
        }

        Chip8 {
            memory: memory,
            regs: regs,
            addressReg: self.addressReg,

            pc: self.pc,
            stack: self.stack.clone(),

            delay_timer: self.delay_timer,
            sound_timer: self.sound_timer,

            screen: screen,
        }
    }
}

impl Chip8 {
    pub fn new() -> Chip8 {
        let mut chip8 = Chip8 {
            memory: [0; 4096],
            regs: [0; 16],
            addressReg: 0,

            pc: PROGRAM_START,
            stack: Vec::new(),

            delay_timer: 0,
            sound_timer: 0,

            screen: [[false; 64]; 32],
        };

        chip8.inject_fontmap();
        chip8
    }

    pub fn inject_fontmap(&mut self) {
        for (offset, byte) in FONTMAP.iter().enumerate() {
            self.memory[FONT_START as usize + offset] = *byte;
        }
    }

    pub fn load_program<R: io::Read>(&mut self, mut program: R) -> io::Result<()> {
        let mut bytes = Vec::new();
        try!(program.read_to_end(&mut bytes));

        for (offset, byte) in bytes.iter().enumerate() {
            self.memory[offset + PROGRAM_START as usize] = *byte;
        }

        Ok(())
    }

    pub fn cycle(&mut self, keys: [bool; 16]) -> Result<(), RuntimeError> {
        use self::RuntimeError::*;

        Chip8::decrement_timer(&mut self.delay_timer);
        Chip8::decrement_timer(&mut self.sound_timer);

        let pc_index = self.pc as usize;
        let opcode_bytes = (self.memory[pc_index] as u16) << 8 | (self.memory[pc_index + 1] as u16);

        let opcode = match Opcode::from_u16(opcode_bytes) {
            Ok(opcode) => opcode,
            Err(err) => return Err(OpcodeErr(err)),
        };
        println!("{:X}: {:?}", opcode_bytes, opcode);

        self.pc += 2;
        try!(self.execute_opcode(opcode, keys));

        Ok(())
    }

    pub fn clear_screen(&mut self) {
        for row in self.screen.iter_mut() {
            for col in row.iter_mut() {
                *col = false;
            }
        }
    }

    // Wrapping is performed in this function, no need to perform it outside
    // Returns if a pixel was unset
    pub fn set_pixel(&mut self, x: usize, y: usize) -> bool {
        // Equivalent to using the mod operator, but faster
        let x = x & 63;
        let y = y & 31;

        let previous_state = self.screen[y][x];
        self.screen[y][x] = !self.screen[y][x];

        // return whether a pixel was previously set and then now unset
        previous_state
    }

    pub fn execute_opcode(&mut self, opcode: Opcode, keys: [bool; 16]) -> Result<(), RuntimeError> {
        use self::RuntimeError::*;
        use opcode::Opcode::*;

        match opcode {
            ClearScreen => self.clear_screen(),
            Return => {
                self.pc = match self.stack.pop() {
                    Some(addr) => addr,
                    None => return Err(EmptyCallStack),
                };
            },

            JumpTo { addr, plus_v0 } => {
                self.pc = addr;

                if plus_v0 {
                    self.pc += self.regs[0] as u16;
                }
            },
            Call(addr) => {
                self.stack.push(self.pc);
                self.pc = addr;
            },

            SkipIfRegEqualConst { not_equal, reg, value } => {
                let mut should_jump = self.regs[reg as usize] == value;

                if not_equal {
                    should_jump = !should_jump; // Effectively computes self.regs[reg] != value
                }

                if should_jump {
                    self.pc += 2;
                }
            },
            SkipIfRegsEqual { not_equal, regs: (v_x, v_y) } => {
                let mut should_jump = self.regs[v_x as usize] == self.regs[v_y as usize];

                if not_equal {
                    should_jump = !should_jump;
                }

                if should_jump {
                    self.pc += 2;
                }
            },

            SetRegToConst { add, reg, value } => {
                if add {
                    let mut value = (self.regs[reg as usize] as u32 + value as u32) & 255;
                    self.regs[reg as usize] = value as u8;
                } else {
                    self.regs[reg as usize] = value;
                }
            },
            SetRegToReg { regs: (v_x, v_y), mode } => {
                let v_x = v_x as usize;
                let v_y = v_y as usize;

                match mode {
                    SetRegMode::Copy => self.regs[v_x] = self.regs[v_y],

                    SetRegMode::Or => self.regs[v_x] |= self.regs[v_y],
                    SetRegMode::And => self.regs[v_x] &= self.regs[v_y],
                    SetRegMode::Xor => self.regs[v_x] ^= self.regs[v_y],

                    SetRegMode::Add => {
                        self.regs[0xF] = 0;

                        let mut reg_value = self.regs[v_x] as usize + self.regs[v_y] as usize;
                        if reg_value > 256 {
                            reg_value -= 256;
                            self.regs[0xF] = 1;
                        }

                        self.regs[v_x] = reg_value as u8;
                    },
                    SetRegMode::Subtract | SetRegMode::InverseSubtract => {
                        self.regs[0xF] = 1;

                        let mut reg_value = if mode == SetRegMode::Subtract {
                            self.regs[v_x] as isize - self.regs[v_y] as isize
                        } else { // Must be InverseSubtract
                            self.regs[v_y] as isize - self.regs[v_x] as isize
                        };

                        if reg_value < 0 {
                            reg_value += 256;
                            self.regs[0xF] = 0;
                        }

                        self.regs[v_x] = reg_value as u8;
                    },

                    // v_y is ignored for the shift opcodes, not sure why
                    SetRegMode::ShiftLeft => {
                        self.regs[0xF] = self.regs[v_x] & 128;

                        self.regs[v_x] <<= 1;
                    },
                    SetRegMode::ShiftRight => {
                        self.regs[0xF] = self.regs[v_x] & 0x1;
                        self.regs[v_x] >>= 1;
                    }
                }
            },

            SetAddressReg(addr) => self.addressReg = addr,
            SetRegToRandom { reg, mask } => {
                let rand: u8 = thread_rng().gen();
                self.regs[reg as usize] = rand & mask;
            },

            DrawSprite { regs: (v_x, v_y), rows } => {
                let x = self.regs[v_x as usize] as usize;
                let y = self.regs[v_y as usize] as usize;

                self.regs[0xF] = 0;

                for row in 0..rows {
                    let sprite_slice = self.memory[(self.addressReg + row as u16) as usize];

                    for col in 0..8 {
                        if (sprite_slice & (128 >> col)) != 0 {
                            if self.set_pixel(x + col as usize, y + row as usize) {
                                self.regs[0xF] = 1;
                            }
                        }
                    }
                }
            },

            SetRegToDelayTimer(reg) => self.regs[reg as usize] = self.delay_timer as u8,

            SetDelayTimerToReg(reg) => self.delay_timer = self.regs[reg as usize] as u16,
            SetSoundTimerToReg(reg) => self.sound_timer = self.regs[reg as usize] as u16,

            AddRegToAddressReg(reg) => self.addressReg += self.regs[reg as usize] as u16,
            SetAddressRegToCharInReg(reg) => {
                let ch = self.regs[reg as usize];
                self.addressReg = FONT_START + ch as u16 * 5;
            },

            WaitForKeyInReg(reg) => {
                // TODO
            },
            SkipIfKeyInRegPressed { not_pressed, reg } => {
                let mut should_jump = keys[self.regs[reg as usize] as usize];

                if not_pressed {
                    should_jump = !should_jump;
                }

                if should_jump {
                    self.pc += 2;
                }
            },

            // See http://en.wikipedia.org/wiki/Binary-coded_decimal
            // n mod 10 => Gets the ones digit out of a number
            RegToBCD(reg) => {
                let number = self.regs[reg as usize];

                let hundreds_digit = number / 100;
                let tens_digit = (number / 10) % 10; // Dividing by ten slides the tens digit into the ones digit
                let ones_digit = number % 10;

                self.memory[(self.addressReg) as usize] = hundreds_digit;
                self.memory[(self.addressReg + 1) as usize] = tens_digit;
                self.memory[(self.addressReg + 2) as usize] = ones_digit;
            },

            DumpRegsToAddr(reg) => {
                for cur_reg in 0..(reg + 1) {
                    self.memory[(self.addressReg + cur_reg as u16)  as usize] = self.regs[cur_reg as usize];
                }
            },
            LoadRegsFromAddr(reg) => {
                for cur_reg in 0..(reg + 1) {
                    self.regs[cur_reg as usize] = self.memory[(self.addressReg + cur_reg as u16) as usize];
                }
            }
        }

        Ok(())
    }

    fn decrement_timer(timer: &mut u16) {
        if *timer != 0 {
            *timer -= 1;
        }
    }
}

impl fmt::Debug for Chip8 {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        try!(writeln!(fmt, "Program Counter: 0x{:X} ({})", self.pc, self.pc));
        try!(writeln!(fmt, "Address Register: 0x{:X} ({})", self.addressReg, self.addressReg));
        try!(writeln!(fmt, "Stack: {:?}", self.stack));
        try!(writeln!(fmt, "Delay Timer: {}", self.delay_timer));
        try!(writeln!(fmt, "Sound Timer: {}", self.sound_timer));

        write!(fmt, "Register Contents: {:?}", self.regs)
    }
}
<|endoftext|>


<commit_before>use crate::error::ParseMessageError;
use crate::structs;
use std::fmt;
use std::str::FromStr;

#[derive(Debug)]
pub enum Pw {
    Uid(u64),
    Nam(String),
    Ent(Ent),
}

#[derive(Debug)]
pub enum Sp {
    Nam(String),
    Ent(Ent),
}

#[derive(Debug)]
pub enum Gr {
    Gid(u64),
    Nam(String),
    Ent(Ent),
}

#[derive(Debug)]
pub enum Ent {
    Set(u32),
    Get(u32),
    End(u32),
}

#[derive(Debug)]
pub enum ClientMessage {
    Key { user: String },
    Pam { user: String },
    CleanUp,
    RateLimit,
    SectorGroups,
    Pw(Pw),
    Sp(Sp),
    Gr(Gr),
}

#[allow(unused)]
#[derive(Debug)]
pub enum DaemonMessage {
    Success,
    Error {
        message: String,
    },
    Key {
        keys: String,
    },
    Pam {
        result: bool,
    },
    RateLimit {
        limit: usize,
    },
    SectorGroups {
        sectors: Vec<structs::SectorGroup>,
    },
    Pw {
        login: String,
        uid: u64,
        gid: u64,
        home: String,
        sh: String,
    },
    Sp {
        login: String,
        pass: String,
    },
    Gr {
        sector: structs::SectorGroup,
    },
}

impl fmt::Display for Ent {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Ent::Set(pid) => write!(f, "set|{}", pid),
            Ent::Get(pid) => write!(f, "get|{}", pid),
            Ent::End(pid) => write!(f, "end|{}", pid),
        }
    }
}

impl fmt::Display for Pw {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Pw::Uid(uid) => write!(f, "uid={}", uid),
            Pw::Nam(name) => write!(f, "name={}", name),
            Pw::Ent(ent) => write!(f, "ent={}", ent),
        }
    }
}

impl fmt::Display for Sp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Sp::Nam(name) => write!(f, "name={}", name),
            Sp::Ent(ent) => write!(f, "ent={}", ent),
        }
    }
}

impl fmt::Display for Gr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Gr::Gid(gid) => write!(f, "gid={}", gid),
            Gr::Nam(name) => write!(f, "name={}", name),
            Gr::Ent(ent) => write!(f, "ent={}", ent),
        }
    }
}

impl fmt::Display for ClientMessage {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ClientMessage::Key { user } => write!(f, "c:key:{}", user),
            ClientMessage::Pam { user } => write!(f, "c:pam:{}", user),
            ClientMessage::CleanUp => write!(f, "c:cleanup"),
            ClientMessage::RateLimit => write!(f, "c:ratelimit"),
            ClientMessage::SectorGroups => write!(f, "c:sectors"),
            ClientMessage::Pw(pw) => write!(f, "c:pw:{}", pw),
            ClientMessage::Sp(sp) => write!(f, "c:sp:{}", sp),
            ClientMessage::Gr(gr) => write!(f, "c:gr:{}", gr),
        }
    }
}

impl fmt::Display for DaemonMessage {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DaemonMessage::Error { message } => write!(f, "d:error:{}", message),
            DaemonMessage::Success => write!(f, "d:success"),
            DaemonMessage::Key { keys } => write!(f, "d:key:{}", keys),
            DaemonMessage::Pam { result } => write!(f, "d:pam:{}", result),
            DaemonMessage::RateLimit { limit } => write!(f, "d:ratelimit:{}", limit),
            DaemonMessage::SectorGroups { sectors } => {
                let ss: Vec<String> = sectors.iter().map(|s| s.to_string()).collect();
                write!(f, "d:sectors:{}", ss.join("\n"))
            }
            DaemonMessage::Pw { login,
                                uid,
                                gid,
                                home,
                                sh, } => write!(f, "d:pw:{}:{}:{}:{}:{}", login, uid, gid, home, sh),
            DaemonMessage::Sp { login, pass } => write!(f, "d:sp:{}:{}", login, pass),
            DaemonMessage::Gr { sector } => write!(f, "d:gr:{}", sector),
        }
    }
}

impl FromStr for Ent {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("set|") {
            Ok(Ent::Set(s.get(4..).unwrap_or_default().parse::<u32>().unwrap()))
        } else if s.starts_with("get|") {
            Ok(Ent::Get(s.get(4..).unwrap_or_default().parse::<u32>().unwrap()))
        } else if s.starts_with("end|") {
            Ok(Ent::End(s.get(4..).unwrap_or_default().parse::<u32>().unwrap()))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for Pw {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("uid=") {
            Ok(Pw::Uid(s.get(4..).unwrap_or_default().parse::<u64>().unwrap()))
        } else if s.starts_with("name=") {
            Ok(Pw::Nam(String::from(s.get(5..).unwrap_or_default())))
        } else if s.starts_with("ent=") {
            Ok(Pw::Ent(s.get(4..).unwrap_or_default().parse::<Ent>().unwrap()))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for Sp {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("name=") {
            Ok(Sp::Nam(String::from(s.get(5..).unwrap_or_default())))
        } else if s.starts_with("ent=") {
            Ok(Sp::Ent(s.get(4..).unwrap_or_default().parse::<Ent>().unwrap()))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for Gr {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("gid=") {
            Ok(Gr::Gid(s.get(4..).unwrap_or_default().parse::<u64>().unwrap()))
        } else if s.starts_with("name=") {
            Ok(Gr::Nam(String::from(s.get(5..).unwrap_or_default())))
        } else if s.starts_with("ent=") {
            Ok(Gr::Ent(s.get(4..).unwrap_or_default().parse::<Ent>().unwrap()))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for ClientMessage {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("c:key:") {
            Ok(ClientMessage::Key { user: String::from(s.get(6..).unwrap_or_default()) })
        } else if s.starts_with("c:pam:") {
            Ok(ClientMessage::Pam { user: String::from(s.get(6..).unwrap_or_default()) })
        } else if s == "c:cleanup" {
            Ok(ClientMessage::CleanUp)
        } else if s == "c:ratelimit" {
            Ok(ClientMessage::RateLimit)
        } else if s == "c:sectors" {
            Ok(ClientMessage::SectorGroups)
        } else if s.starts_with("c:pw:") {
            Ok(ClientMessage::Pw(s.get(5..).unwrap_or_default().parse::<Pw>()?))
        } else if s.starts_with("c:sp:") {
            Ok(ClientMessage::Sp(s.get(5..).unwrap_or_default().parse::<Sp>()?))
        } else if s.starts_with("c:gr:") {
            Ok(ClientMessage::Gr(s.get(5..).unwrap_or_default().parse::<Gr>()?))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for DaemonMessage {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("d:key:") {
            Ok(DaemonMessage::Key { keys: String::from(s.get(6..).unwrap_or_default()) })
        } else if s.starts_with("d:pam:") {
            Ok(DaemonMessage::Pam { result: FromStr::from_str(s.get(6..).unwrap_or("false")).unwrap_or(false) })
        } else if s == "d:success" {
            Ok(DaemonMessage::Success)
        } else if s.starts_with("d:ratelimit:") {
            Ok(DaemonMessage::RateLimit { limit: s.get(12..).unwrap_or("0").parse::<usize>().unwrap_or(0) })
        } else if s.starts_with("d:sectors:") {
            let sectors = s.get(10..)
                           .unwrap_or_default()
                           .split('\n')
                           .filter_map(|l| l.parse::<structs::SectorGroup>().ok())
                           .collect();
            Ok(DaemonMessage::SectorGroups { sectors })
        } else if s.starts_with("d:pw:") {
            let fields: Vec<String> = s.get(5..)
                                       .unwrap_or_default()
                                       .split(':')
                                       .map(|s| s.to_string())
                                       .collect();
            if fields.len() < 5 {
                return Err(ParseMessageError::ParseDaemonMessageError);
            }
            let login: String = fields[0].clone();
            let home: String = fields[3].clone();
            let sh: String = fields[4].clone();
            match (fields[1].parse::<u64>(), fields[2].parse::<u64>()) {
                (Ok(uid), Ok(gid)) => Ok(DaemonMessage::Pw { login,
                                                             uid,
                                                             gid,
                                                             home,
                                                             sh }),
                _ => Err(ParseMessageError::ParseDaemonMessageError),
            }
        } else if s.starts_with("d:sp:") {
            let fields: Vec<String> = s.get(5..)
                                       .unwrap_or_default()
                                       .split(':')
                                       .map(|s| s.to_string())
                                       .collect();
            if fields.len() < 2 {
                return Err(ParseMessageError::ParseDaemonMessageError);
            }
            Ok(DaemonMessage::Sp { login: fields[0].clone(),
                                   pass: fields[1].clone() })
        } else if s.starts_with("d:gr:") {
            match s.get(5..).unwrap_or_default().parse::<structs::SectorGroup>() {
                Ok(sector) => Ok(DaemonMessage::Gr { sector }),
                _ => Err(ParseMessageError::ParseDaemonMessageError),
            }
        } else {
            Err(ParseMessageError::ParseDaemonMessageError)
        }
    }
}
<commit_msg>Tweak pragma style<commit_after>use crate::error::ParseMessageError;
use crate::structs;
use std::fmt;
use std::str::FromStr;

#[derive(Debug)]
pub enum Pw {
    Uid(u64),
    Nam(String),
    Ent(Ent),
}

#[derive(Debug)]
pub enum Sp {
    Nam(String),
    Ent(Ent),
}

#[derive(Debug)]
pub enum Gr {
    Gid(u64),
    Nam(String),
    Ent(Ent),
}

#[derive(Debug)]
pub enum Ent {
    Set(u32),
    Get(u32),
    End(u32),
}

#[derive(Debug)]
pub enum ClientMessage {
    Key { user: String },
    Pam { user: String },
    CleanUp,
    RateLimit,
    SectorGroups,
    Pw(Pw),
    Sp(Sp),
    Gr(Gr),
}

#[allow(dead_code)]
#[derive(Debug)]
pub enum DaemonMessage {
    Success,
    Error {
        message: String,
    },
    Key {
        keys: String,
    },
    Pam {
        result: bool,
    },
    RateLimit {
        limit: usize,
    },
    SectorGroups {
        sectors: Vec<structs::SectorGroup>,
    },
    Pw {
        login: String,
        uid: u64,
        gid: u64,
        home: String,
        sh: String,
    },
    Sp {
        login: String,
        pass: String,
    },
    Gr {
        sector: structs::SectorGroup,
    },
}

impl fmt::Display for Ent {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Ent::Set(pid) => write!(f, "set|{}", pid),
            Ent::Get(pid) => write!(f, "get|{}", pid),
            Ent::End(pid) => write!(f, "end|{}", pid),
        }
    }
}

impl fmt::Display for Pw {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Pw::Uid(uid) => write!(f, "uid={}", uid),
            Pw::Nam(name) => write!(f, "name={}", name),
            Pw::Ent(ent) => write!(f, "ent={}", ent),
        }
    }
}

impl fmt::Display for Sp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Sp::Nam(name) => write!(f, "name={}", name),
            Sp::Ent(ent) => write!(f, "ent={}", ent),
        }
    }
}

impl fmt::Display for Gr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Gr::Gid(gid) => write!(f, "gid={}", gid),
            Gr::Nam(name) => write!(f, "name={}", name),
            Gr::Ent(ent) => write!(f, "ent={}", ent),
        }
    }
}

impl fmt::Display for ClientMessage {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ClientMessage::Key { user } => write!(f, "c:key:{}", user),
            ClientMessage::Pam { user } => write!(f, "c:pam:{}", user),
            ClientMessage::CleanUp => write!(f, "c:cleanup"),
            ClientMessage::RateLimit => write!(f, "c:ratelimit"),
            ClientMessage::SectorGroups => write!(f, "c:sectors"),
            ClientMessage::Pw(pw) => write!(f, "c:pw:{}", pw),
            ClientMessage::Sp(sp) => write!(f, "c:sp:{}", sp),
            ClientMessage::Gr(gr) => write!(f, "c:gr:{}", gr),
        }
    }
}

impl fmt::Display for DaemonMessage {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DaemonMessage::Error { message } => write!(f, "d:error:{}", message),
            DaemonMessage::Success => write!(f, "d:success"),
            DaemonMessage::Key { keys } => write!(f, "d:key:{}", keys),
            DaemonMessage::Pam { result } => write!(f, "d:pam:{}", result),
            DaemonMessage::RateLimit { limit } => write!(f, "d:ratelimit:{}", limit),
            DaemonMessage::SectorGroups { sectors } => {
                let ss: Vec<String> = sectors.iter().map(|s| s.to_string()).collect();
                write!(f, "d:sectors:{}", ss.join("\n"))
            }
            DaemonMessage::Pw { login,
                                uid,
                                gid,
                                home,
                                sh, } => write!(f, "d:pw:{}:{}:{}:{}:{}", login, uid, gid, home, sh),
            DaemonMessage::Sp { login, pass } => write!(f, "d:sp:{}:{}", login, pass),
            DaemonMessage::Gr { sector } => write!(f, "d:gr:{}", sector),
        }
    }
}

impl FromStr for Ent {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("set|") {
            Ok(Ent::Set(s.get(4..).unwrap_or_default().parse::<u32>().unwrap()))
        } else if s.starts_with("get|") {
            Ok(Ent::Get(s.get(4..).unwrap_or_default().parse::<u32>().unwrap()))
        } else if s.starts_with("end|") {
            Ok(Ent::End(s.get(4..).unwrap_or_default().parse::<u32>().unwrap()))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for Pw {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("uid=") {
            Ok(Pw::Uid(s.get(4..).unwrap_or_default().parse::<u64>().unwrap()))
        } else if s.starts_with("name=") {
            Ok(Pw::Nam(String::from(s.get(5..).unwrap_or_default())))
        } else if s.starts_with("ent=") {
            Ok(Pw::Ent(s.get(4..).unwrap_or_default().parse::<Ent>().unwrap()))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for Sp {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("name=") {
            Ok(Sp::Nam(String::from(s.get(5..).unwrap_or_default())))
        } else if s.starts_with("ent=") {
            Ok(Sp::Ent(s.get(4..).unwrap_or_default().parse::<Ent>().unwrap()))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for Gr {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("gid=") {
            Ok(Gr::Gid(s.get(4..).unwrap_or_default().parse::<u64>().unwrap()))
        } else if s.starts_with("name=") {
            Ok(Gr::Nam(String::from(s.get(5..).unwrap_or_default())))
        } else if s.starts_with("ent=") {
            Ok(Gr::Ent(s.get(4..).unwrap_or_default().parse::<Ent>().unwrap()))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for ClientMessage {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("c:key:") {
            Ok(ClientMessage::Key { user: String::from(s.get(6..).unwrap_or_default()) })
        } else if s.starts_with("c:pam:") {
            Ok(ClientMessage::Pam { user: String::from(s.get(6..).unwrap_or_default()) })
        } else if s == "c:cleanup" {
            Ok(ClientMessage::CleanUp)
        } else if s == "c:ratelimit" {
            Ok(ClientMessage::RateLimit)
        } else if s == "c:sectors" {
            Ok(ClientMessage::SectorGroups)
        } else if s.starts_with("c:pw:") {
            Ok(ClientMessage::Pw(s.get(5..).unwrap_or_default().parse::<Pw>()?))
        } else if s.starts_with("c:sp:") {
            Ok(ClientMessage::Sp(s.get(5..).unwrap_or_default().parse::<Sp>()?))
        } else if s.starts_with("c:gr:") {
            Ok(ClientMessage::Gr(s.get(5..).unwrap_or_default().parse::<Gr>()?))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for DaemonMessage {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("d:key:") {
            Ok(DaemonMessage::Key { keys: String::from(s.get(6..).unwrap_or_default()) })
        } else if s.starts_with("d:pam:") {
            Ok(DaemonMessage::Pam { result: FromStr::from_str(s.get(6..).unwrap_or("false")).unwrap_or(false) })
        } else if s == "d:success" {
            Ok(DaemonMessage::Success)
        } else if s.starts_with("d:ratelimit:") {
            Ok(DaemonMessage::RateLimit { limit: s.get(12..).unwrap_or("0").parse::<usize>().unwrap_or(0) })
        } else if s.starts_with("d:sectors:") {
            let sectors = s.get(10..)
                           .unwrap_or_default()
                           .split('\n')
                           .filter_map(|l| l.parse::<structs::SectorGroup>().ok())
                           .collect();
            Ok(DaemonMessage::SectorGroups { sectors })
        } else if s.starts_with("d:pw:") {
            let fields: Vec<String> = s.get(5..)
                                       .unwrap_or_default()
                                       .split(':')
                                       .map(|s| s.to_string())
                                       .collect();
            if fields.len() < 5 {
                return Err(ParseMessageError::ParseDaemonMessageError);
            }
            let login: String = fields[0].clone();
            let home: String = fields[3].clone();
            let sh: String = fields[4].clone();
            match (fields[1].parse::<u64>(), fields[2].parse::<u64>()) {
                (Ok(uid), Ok(gid)) => Ok(DaemonMessage::Pw { login,
                                                             uid,
                                                             gid,
                                                             home,
                                                             sh }),
                _ => Err(ParseMessageError::ParseDaemonMessageError),
            }
        } else if s.starts_with("d:sp:") {
            let fields: Vec<String> = s.get(5..)
                                       .unwrap_or_default()
                                       .split(':')
                                       .map(|s| s.to_string())
                                       .collect();
            if fields.len() < 2 {
                return Err(ParseMessageError::ParseDaemonMessageError);
            }
            Ok(DaemonMessage::Sp { login: fields[0].clone(),
                                   pass: fields[1].clone() })
        } else if s.starts_with("d:gr:") {
            match s.get(5..).unwrap_or_default().parse::<structs::SectorGroup>() {
                Ok(sector) => Ok(DaemonMessage::Gr { sector }),
                _ => Err(ParseMessageError::ParseDaemonMessageError),
            }
        } else {
            Err(ParseMessageError::ParseDaemonMessageError)
        }
    }
}
<|endoftext|>


<commit_before>use std::string::{String};
use std::fmt;

#[allow(dead_code)]
pub type IrcProtocolMessage = self::IrcProtocolMessage::IrcProtocolMessage;
pub mod IrcProtocolMessage {
    #[deriving(Clone)]
    pub enum IrcProtocolMessage {
        Ping(String, Option<String>),
        Pong(String),
        Notice(String),
        Join(String),
        Numeric(u16, Vec<String>),
        // parsed but not processed into a safe message type. command, rest
        Unknown(String, Vec<String>)
    }
}

impl fmt::Show for IrcProtocolMessage {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            IrcProtocolMessage::Ping(ref s1, None) => write!(f, "PING {}", s1),
            IrcProtocolMessage::Ping(ref s1, Some(ref s2)) => write!(f, "PING {} {}", s1, s2),
            IrcProtocolMessage::Pong(ref s1) => write!(f, "PONG {}", s1),
            IrcProtocolMessage::Notice(ref s1) => write!(f, "NOTICE {}", s1),
            _ => write!(f, "WHAT")
        }
    }
}

#[deriving(PartialEq, Clone)]
pub struct IrcHostmask {
    nick: String,
    user: String,
    host: String
}

impl fmt::Show for IrcHostmask {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}!{}@{}",
            self.nick.as_slice(),
            self.user.as_slice(),
            self.host.as_slice())
    }
}

#[deriving(PartialEq, Clone, Show)]
pub enum IrcPrefix {
    IrcHostmaskPrefix(IrcHostmask),
    IrcOtherPrefix(String)
}

#[deriving(Clone)]
pub struct IrcMessage {
    prefix: Option<IrcPrefix>,
    prefix_raw: Option<String>,
    message: IrcProtocolMessage,
    command: String,
    args: Vec<String>
}

fn parse_prefix(text: &str) -> IrcPrefix {
    let parts: Vec<&str> = text.splitn('!', 1).collect();
    let (nick, rest) = match parts.as_slice() {
        [_] => return IrcOtherPrefix(String::from_str(text)),
        [nick, rest] => (nick, rest),
        _ => fail!("programmer error")
    };

    let parts: Vec<&str> = rest.splitn('@', 1).collect();
    let (user, host) = match parts.as_slice() {
        [_] => return IrcOtherPrefix(String::from_str(text)),
        [user, rest] => (user, rest),
        _ => fail!("programmer error")
    };
    IrcHostmaskPrefix(IrcHostmask {
        nick: String::from_str(nick),
        user: String::from_str(user),
        host: String::from_str(host),
    })
}

fn parse_message_args(text: &str) -> Result<Vec<String>, Option<String>> {
    if text.len() == 0 {
        return Err(from_str("Invalid IRC message"));
    }
    if text.char_at(0) == ':' {
        return Ok(vec![String::from_str(text.slice_from(1))]);
    }

    let (arg_parts, rest) = match text.find_str(" :") {
        Some(val) => (text.slice_to(val), Some(text.slice_from(val + 2))),
        None => (text, None)
    };

    let mut output: Vec<String> = arg_parts.split(' ')
            .map(|s| String::from_str(s)).collect();

    if rest.is_some() {
        output.push(String::from_str(rest.unwrap()));
    }
    Ok(output)
}

fn parse_message_rest(text: &str) -> Result<(String, Vec<String>), Option<String>> {
    let parts: Vec<&str> = text.splitn(' ', 1).collect();
    let args = match parse_message_args(parts[1]) {
        Ok(args) => args,
        Err(err) => return Err(err)
    };
    Ok((String::from_str(parts[0]), args))
}

impl IrcMessage {
    pub fn from_str(text: &str) -> Result<IrcMessage, Option<String>> {
        if text.len() == 0 {
            return Err(from_str("Invalid IRC message"));
        }
        let (prefix, command, mut args) = if text.char_at(0) == ':' {
                let parts: Vec<&str> = text.splitn(' ', 1).collect();
                if parts.len() < 2 {
                    return Err(from_str("Invalid IRC message"));
                }
                let (command, args) = match parse_message_rest(parts[1]) {
                    Ok(result) => result,
                    Err(err) => return Err(Some(format!("Invalid IRC message: {}", err)))
                };

                (Some(String::from_str(parts[0].slice_from(1))), command, args)
            } else {
                assert!(text.len() > 0);
                let (command, args) = match parse_message_rest(text) {
                    Ok(result) => result,
                    Err(err) => return Err(Some(format!("Invalid IRC message: {}", err)))
                };
                (None, command, args)
            };

        let message_command = command.clone();
        let message_args = args.clone();

        let message = match (command.as_slice(), args.len()) {
            ("PING", 1..2) => {
                IrcProtocolMessage::Ping(args.remove(0).unwrap(), args.remove(0))
            },
            ("PING", _) => return Err(from_str(
                "Invalid IRC message: too many arguments to PING")),
            ("PONG", 1) => IrcProtocolMessage::Pong(args.remove(0).unwrap()),
            ("PONG", _) => return Err(from_str(
                "Invalid IRC message: too many arguments to PONG")),
            (_, _) => {
                match from_str(command.as_slice()) {
                    Some(num) => IrcProtocolMessage::Numeric(num, args),
                    None => IrcProtocolMessage::Unknown(command, args)
                }
            }
        };

        let prefix_parsed = match prefix {
            Some(ref pref) => Some(parse_prefix(pref.as_slice())),
            None => None
        };

        Ok(IrcMessage {
            prefix: prefix_parsed,
            prefix_raw: prefix,
            message: message,
            command: message_command,
            args: message_args
        })
    }

    pub fn source_nick(&self) -> Option<String> {
        match self.prefix {
            Some(IrcHostmaskPrefix(ref hostmask)) => Some(hostmask.nick.clone()),
            Some(IrcOtherPrefix(_)) => None,
            None => None
        }
    }

    pub fn get_prefix<'a>(&'a self) -> Option<&'a IrcPrefix> {
        match self.prefix {
            Some(ref pref) => Some(pref),
            None => None
        }
    }

    pub fn get_message<'a>(&'a self) -> &'a IrcProtocolMessage {
        &self.message
    }

    pub fn get_command<'a>(&'a self) -> &'a String {
        &self.command
    }

    pub fn get_args(&self) -> &Vec<String> {
        &self.args
    }

    pub fn get_arg<'a>(&'a self, i: uint) -> &'a String {
        &self.args[i]
    }
}

impl fmt::Show for IrcMessage {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut arg_string = String::new();
        arg_string = arg_string.append("[");
        for part in self.args.iter().map(|s| s.as_slice()) {
            arg_string = arg_string.append(format!("{}, ", part.as_slice()).as_slice());
        }
        arg_string = arg_string.append("]");

        match self.prefix_raw {
            Some(ref prefix) => write!(f, "IrcMessage({}, {}, {})",
                prefix.as_slice(), self.command.as_slice(), arg_string.as_slice()),
            None => write!(f, "IrcMessage({}, {})",
                self.command.as_slice(), arg_string.as_slice())
        }
    }
}

#[test]
fn test_irc_message_general() {
    match IrcMessage::from_str("") {
        Ok(_) => {
            fail!("empty string is invalid")
        },
        Err(_) => ()
    };

    match IrcMessage::from_str(":") {
        Ok(_) => {
            fail!("single colon is invalid")
        },
        Err(_) => ()
    };

    match IrcMessage::from_str(" ") {
        Ok(_) => {
            fail!("single space is invalid")
        },
        Err(_) => ()
    };

    match IrcMessage::from_str(":nick!user@host PING server1 :server2") {
        Ok(message) => {
            match message.prefix {
                Some(IrcHostmaskPrefix(ref data)) => {
                    assert_eq!(data.nick.as_slice(), "nick");
                    assert_eq!(data.user.as_slice(), "user");
                    assert_eq!(data.host.as_slice(), "host");
                }
                _ => fail!("invalid parsed prefix")
            };
        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str("PING server1") {
        Ok(message) => {
            assert_eq!(message.prefix_raw, None);
            assert_eq!(message.command.as_slice(), "PING");
            assert_eq!(message.args.len(), 1);
        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str("PING server1 server2") {
        Ok(message) => {
            assert_eq!(message.prefix_raw, None);
            assert_eq!(message.command.as_slice(), "PING");
            assert_eq!(message.args.len(), 2);
        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str("PING server1 server2 server3") {
        Ok(_) => fail!("should fail to parse"),
        Err(_) => ()
    };

    match IrcMessage::from_str(":somewhere PING server1") {
        Ok(message) => {
            assert_eq!(message.prefix_raw, Some(String::from_str("somewhere")));
            assert_eq!(message.command.as_slice(), "PING");
            assert_eq!(message.args.len(), 1);
        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str(":somewhere PING server1 server2") {
        Ok(message) => {
            assert_eq!(message.prefix_raw, Some(String::from_str("somewhere")));
            assert_eq!(message.command.as_slice(), "PING");
            assert_eq!(message.args.len(), 2);
            assert_eq!(message.args[0].as_slice(), "server1");
            assert_eq!(message.args[1].as_slice(), "server2");
            match message.message {
                IrcProtocolMessage::Ping(s1, s2) => {
                    assert_eq!(s1, String::from_str("server1"));
                    assert_eq!(s2, Some(String::from_str("server2")));
                },
                _ => assert!(false)
            };

        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str(":somewhere PING server1 :server2") {
        Ok(message) => {
            assert_eq!(message.prefix, Some(IrcOtherPrefix(String::from_str("somewhere"))));
            assert_eq!(message.prefix_raw, Some(String::from_str("somewhere")));
            assert_eq!(message.command.as_slice(), "PING");
            assert_eq!(message.args.len(), 2);
            assert_eq!(message.args[0].as_slice(), "server1");
            assert_eq!(message.args[1].as_slice(), "server2");
        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str(":somewhere PING :server1 server2") {
        Ok(message) => {
            assert_eq!(message.prefix_raw, Some(String::from_str("somewhere")));
            assert_eq!(message.command.as_slice(), "PING");
            assert_eq!(message.args.len(), 1);
            assert_eq!(message.args[0].as_slice(), "server1 server2");
        },
        Err(_) => fail!("failed to parse")
    };
}

#[test]
fn test_irc_message_numerics() {
    match IrcMessage::from_str(":somewhere 001 nick :blah blah") {
        Ok(message) => {
            assert_eq!(message.prefix_raw, Some(String::from_str("somewhere")));
            assert_eq!(message.command.as_slice(), "001");
            match message.message {
                IrcProtocolMessage::Numeric(num, _) => assert_eq!(num, 1),
                _ => fail!("numbers should parse as numerics")
            }

        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str("001 nick :blah blah") {
        Ok(message) => {
            assert_eq!(message.prefix_raw, None);
            assert_eq!(message.command.as_slice(), "001");
            match message.message {
                IrcProtocolMessage::Numeric(num, _) => assert_eq!(num, 1),
                _ => fail!("numbers should parse as numerics")
            }

        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str("366 arg") {
        Ok(message) => {
            match message.message {
                IrcProtocolMessage::Numeric(num, _) => assert_eq!(num, 366),
                _ => fail!("numbers should parse as numerics")
            }

        },
        Err(_) => fail!("failed to parse")
    };
}
<commit_msg>Fix recent splitn API break<commit_after>use std::string::{String};
use std::fmt;

#[allow(dead_code)]
pub type IrcProtocolMessage = self::IrcProtocolMessage::IrcProtocolMessage;
pub mod IrcProtocolMessage {
    #[deriving(Clone)]
    pub enum IrcProtocolMessage {
        Ping(String, Option<String>),
        Pong(String),
        Notice(String),
        Join(String),
        Numeric(u16, Vec<String>),
        // parsed but not processed into a safe message type. command, rest
        Unknown(String, Vec<String>)
    }
}

impl fmt::Show for IrcProtocolMessage {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            IrcProtocolMessage::Ping(ref s1, None) => write!(f, "PING {}", s1),
            IrcProtocolMessage::Ping(ref s1, Some(ref s2)) => write!(f, "PING {} {}", s1, s2),
            IrcProtocolMessage::Pong(ref s1) => write!(f, "PONG {}", s1),
            IrcProtocolMessage::Notice(ref s1) => write!(f, "NOTICE {}", s1),
            _ => write!(f, "WHAT")
        }
    }
}

#[deriving(PartialEq, Clone)]
pub struct IrcHostmask {
    nick: String,
    user: String,
    host: String
}

impl fmt::Show for IrcHostmask {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}!{}@{}",
            self.nick.as_slice(),
            self.user.as_slice(),
            self.host.as_slice())
    }
}

#[deriving(PartialEq, Clone, Show)]
pub enum IrcPrefix {
    IrcHostmaskPrefix(IrcHostmask),
    IrcOtherPrefix(String)
}

#[deriving(Clone)]
pub struct IrcMessage {
    prefix: Option<IrcPrefix>,
    prefix_raw: Option<String>,
    message: IrcProtocolMessage,
    command: String,
    args: Vec<String>
}

fn parse_prefix(text: &str) -> IrcPrefix {
    let parts: Vec<&str> = text.splitn(1, '!').collect();
    let (nick, rest) = match parts.as_slice() {
        [_] => return IrcOtherPrefix(String::from_str(text)),
        [nick, rest] => (nick, rest),
        _ => fail!("programmer error")
    };

    let parts: Vec<&str> = rest.splitn(1, '@').collect();
    let (user, host) = match parts.as_slice() {
        [_] => return IrcOtherPrefix(String::from_str(text)),
        [user, rest] => (user, rest),
        _ => fail!("programmer error")
    };
    IrcHostmaskPrefix(IrcHostmask {
        nick: String::from_str(nick),
        user: String::from_str(user),
        host: String::from_str(host),
    })
}

fn parse_message_args(text: &str) -> Result<Vec<String>, Option<String>> {
    if text.len() == 0 {
        return Err(from_str("Invalid IRC message"));
    }
    if text.char_at(0) == ':' {
        return Ok(vec![String::from_str(text.slice_from(1))]);
    }

    let (arg_parts, rest) = match text.find_str(" :") {
        Some(val) => (text.slice_to(val), Some(text.slice_from(val + 2))),
        None => (text, None)
    };

    let mut output: Vec<String> = arg_parts.split(' ')
            .map(|s| String::from_str(s)).collect();

    if rest.is_some() {
        output.push(String::from_str(rest.unwrap()));
    }
    Ok(output)
}

fn parse_message_rest(text: &str) -> Result<(String, Vec<String>), Option<String>> {
    let parts: Vec<&str> = text.splitn(1, ' ').collect();
    let args = match parse_message_args(parts[1]) {
        Ok(args) => args,
        Err(err) => return Err(err)
    };
    Ok((String::from_str(parts[0]), args))
}

impl IrcMessage {
    pub fn from_str(text: &str) -> Result<IrcMessage, Option<String>> {
        if text.len() == 0 {
            return Err(from_str("Invalid IRC message"));
        }
        let (prefix, command, mut args) = if text.char_at(0) == ':' {
                let parts: Vec<&str> = text.splitn(1, ' ').collect();
                if parts.len() < 2 {
                    return Err(from_str("Invalid IRC message"));
                }
                let (command, args) = match parse_message_rest(parts[1]) {
                    Ok(result) => result,
                    Err(err) => return Err(Some(format!("Invalid IRC message: {}", err)))
                };

                (Some(String::from_str(parts[0].slice_from(1))), command, args)
            } else {
                assert!(text.len() > 0);
                let (command, args) = match parse_message_rest(text) {
                    Ok(result) => result,
                    Err(err) => return Err(Some(format!("Invalid IRC message: {}", err)))
                };
                (None, command, args)
            };

        let message_command = command.clone();
        let message_args = args.clone();

        let message = match (command.as_slice(), args.len()) {
            ("PING", 1..2) => {
                IrcProtocolMessage::Ping(args.remove(0).unwrap(), args.remove(0))
            },
            ("PING", _) => return Err(from_str(
                "Invalid IRC message: too many arguments to PING")),
            ("PONG", 1) => IrcProtocolMessage::Pong(args.remove(0).unwrap()),
            ("PONG", _) => return Err(from_str(
                "Invalid IRC message: too many arguments to PONG")),
            (_, _) => {
                match from_str(command.as_slice()) {
                    Some(num) => IrcProtocolMessage::Numeric(num, args),
                    None => IrcProtocolMessage::Unknown(command, args)
                }
            }
        };

        let prefix_parsed = match prefix {
            Some(ref pref) => Some(parse_prefix(pref.as_slice())),
            None => None
        };

        Ok(IrcMessage {
            prefix: prefix_parsed,
            prefix_raw: prefix,
            message: message,
            command: message_command,
            args: message_args
        })
    }

    pub fn source_nick(&self) -> Option<String> {
        match self.prefix {
            Some(IrcHostmaskPrefix(ref hostmask)) => Some(hostmask.nick.clone()),
            Some(IrcOtherPrefix(_)) => None,
            None => None
        }
    }

    pub fn get_prefix<'a>(&'a self) -> Option<&'a IrcPrefix> {
        match self.prefix {
            Some(ref pref) => Some(pref),
            None => None
        }
    }

    pub fn get_message<'a>(&'a self) -> &'a IrcProtocolMessage {
        &self.message
    }

    pub fn get_command<'a>(&'a self) -> &'a String {
        &self.command
    }

    pub fn get_args(&self) -> &Vec<String> {
        &self.args
    }

    pub fn get_arg<'a>(&'a self, i: uint) -> &'a String {
        &self.args[i]
    }
}

impl fmt::Show for IrcMessage {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut arg_string = String::new();
        arg_string = arg_string.append("[");
        for part in self.args.iter().map(|s| s.as_slice()) {
            arg_string = arg_string.append(format!("{}, ", part.as_slice()).as_slice());
        }
        arg_string = arg_string.append("]");

        match self.prefix_raw {
            Some(ref prefix) => write!(f, "IrcMessage({}, {}, {})",
                prefix.as_slice(), self.command.as_slice(), arg_string.as_slice()),
            None => write!(f, "IrcMessage({}, {})",
                self.command.as_slice(), arg_string.as_slice())
        }
    }
}

#[test]
fn test_irc_message_general() {
    match IrcMessage::from_str("") {
        Ok(_) => {
            fail!("empty string is invalid")
        },
        Err(_) => ()
    };

    match IrcMessage::from_str(":") {
        Ok(_) => {
            fail!("single colon is invalid")
        },
        Err(_) => ()
    };

    match IrcMessage::from_str(" ") {
        Ok(_) => {
            fail!("single space is invalid")
        },
        Err(_) => ()
    };

    match IrcMessage::from_str(":nick!user@host PING server1 :server2") {
        Ok(message) => {
            match message.prefix {
                Some(IrcHostmaskPrefix(ref data)) => {
                    assert_eq!(data.nick.as_slice(), "nick");
                    assert_eq!(data.user.as_slice(), "user");
                    assert_eq!(data.host.as_slice(), "host");
                }
                _ => fail!("invalid parsed prefix")
            };
        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str("PING server1") {
        Ok(message) => {
            assert_eq!(message.prefix_raw, None);
            assert_eq!(message.command.as_slice(), "PING");
            assert_eq!(message.args.len(), 1);
        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str("PING server1 server2") {
        Ok(message) => {
            assert_eq!(message.prefix_raw, None);
            assert_eq!(message.command.as_slice(), "PING");
            assert_eq!(message.args.len(), 2);
        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str("PING server1 server2 server3") {
        Ok(_) => fail!("should fail to parse"),
        Err(_) => ()
    };

    match IrcMessage::from_str(":somewhere PING server1") {
        Ok(message) => {
            assert_eq!(message.prefix_raw, Some(String::from_str("somewhere")));
            assert_eq!(message.command.as_slice(), "PING");
            assert_eq!(message.args.len(), 1);
        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str(":somewhere PING server1 server2") {
        Ok(message) => {
            assert_eq!(message.prefix_raw, Some(String::from_str("somewhere")));
            assert_eq!(message.command.as_slice(), "PING");
            assert_eq!(message.args.len(), 2);
            assert_eq!(message.args[0].as_slice(), "server1");
            assert_eq!(message.args[1].as_slice(), "server2");
            match message.message {
                IrcProtocolMessage::Ping(s1, s2) => {
                    assert_eq!(s1, String::from_str("server1"));
                    assert_eq!(s2, Some(String::from_str("server2")));
                },
                _ => assert!(false)
            };

        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str(":somewhere PING server1 :server2") {
        Ok(message) => {
            assert_eq!(message.prefix, Some(IrcOtherPrefix(String::from_str("somewhere"))));
            assert_eq!(message.prefix_raw, Some(String::from_str("somewhere")));
            assert_eq!(message.command.as_slice(), "PING");
            assert_eq!(message.args.len(), 2);
            assert_eq!(message.args[0].as_slice(), "server1");
            assert_eq!(message.args[1].as_slice(), "server2");
        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str(":somewhere PING :server1 server2") {
        Ok(message) => {
            assert_eq!(message.prefix_raw, Some(String::from_str("somewhere")));
            assert_eq!(message.command.as_slice(), "PING");
            assert_eq!(message.args.len(), 1);
            assert_eq!(message.args[0].as_slice(), "server1 server2");
        },
        Err(_) => fail!("failed to parse")
    };
}

#[test]
fn test_irc_message_numerics() {
    match IrcMessage::from_str(":somewhere 001 nick :blah blah") {
        Ok(message) => {
            assert_eq!(message.prefix_raw, Some(String::from_str("somewhere")));
            assert_eq!(message.command.as_slice(), "001");
            match message.message {
                IrcProtocolMessage::Numeric(num, _) => assert_eq!(num, 1),
                _ => fail!("numbers should parse as numerics")
            }

        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str("001 nick :blah blah") {
        Ok(message) => {
            assert_eq!(message.prefix_raw, None);
            assert_eq!(message.command.as_slice(), "001");
            match message.message {
                IrcProtocolMessage::Numeric(num, _) => assert_eq!(num, 1),
                _ => fail!("numbers should parse as numerics")
            }

        },
        Err(_) => fail!("failed to parse")
    };

    match IrcMessage::from_str("366 arg") {
        Ok(message) => {
            match message.message {
                IrcProtocolMessage::Numeric(num, _) => assert_eq!(num, 366),
                _ => fail!("numbers should parse as numerics")
            }

        },
        Err(_) => fail!("failed to parse")
    };
}
<|endoftext|>