<commit_before>use std::path::PathBuf;
use std::default::Default;
use std::collections::BTreeMap;

use quire::validate::{Validator, Structure, Sequence, Scalar, Numeric, Enum};
use quire::validate::{Mapping};
use self::Volume::*;

#[derive(RustcDecodable, RustcEncodable, Clone, PartialEq, Eq)]
pub struct TmpfsInfo {
    pub size: usize,
    pub mode: u32,
}

#[derive(RustcDecodable, RustcEncodable, Clone, PartialEq, Eq)]
pub struct PersistentInfo {
    pub path: PathBuf,
    pub mkdir: bool,
    pub mode: u32,
    pub user: u32,
    pub group: u32,
}

#[derive(RustcDecodable, RustcEncodable, Clone, PartialEq, Eq)]
pub struct StatedirInfo {
    pub path: PathBuf,
    pub mode: u32,
    pub user: u32,
    pub group: u32,
}

#[derive(RustcDecodable, RustcEncodable, Clone, PartialEq, Eq)]
pub enum Volume {
    Readonly(PathBuf),
    Persistent(PersistentInfo),
    Tmpfs(TmpfsInfo),
    Statedir(StatedirInfo),
}

#[derive(RustcDecodable, RustcEncodable, Debug, PartialEq, Eq)]
pub enum ContainerKind {
    Daemon,
    Command,
}

#[derive(RustcDecodable, RustcEncodable)]
pub struct ResolvConf {
    pub copy_from_host: bool,
}

#[derive(RustcDecodable, RustcEncodable)]
pub struct HostsFile {
    pub localhost: bool,
    pub public_hostname: bool,
}

#[derive(RustcDecodable, RustcEncodable, Clone, Copy, Debug)]
pub struct IdMap {
    pub inside: u32,
    pub outside: u32,
    pub count: u32,
}

#[derive(RustcDecodable, RustcEncodable)]
pub struct ContainerConfig {
    pub kind: ContainerKind,
    pub volumes: BTreeMap<String, Volume>,
    pub user_id: u32,
    pub group_id: u32,
    pub restart_timeout: f32,
    pub memory_limit: u64,
    pub fileno_limit: u64,
    pub cpu_shares: usize,
    pub executable: String,
    pub arguments: Vec<String>,
    pub environ: BTreeMap<String, String>,
    pub workdir: PathBuf,
    pub resolv_conf: ResolvConf,
    pub hosts_file: HostsFile,
    pub uid_map: Vec<IdMap>,
    pub gid_map: Vec<IdMap>,
    pub stdout_stderr_file: Option<PathBuf>,
    pub restart_process_only: bool,
}

fn mapping_validator<'x>() -> Box<Validator + 'x> {
    return Box::new(Sequence {
        element: Box::new(Structure { members: vec!(
            ("inside".to_string(), Box::new(Numeric {
                default: None,
                .. Default::default() }) as Box<Validator>),
            ("outside".to_string(), Box::new(Numeric {
                default: None,
                .. Default::default() }) as Box<Validator>),
            ("count".to_string(), Box::new(Numeric {
                default: None,
                .. Default::default() }) as Box<Validator>),
            ), .. Default::default() }) as Box<Validator>,
        .. Default::default() }) as Box<Validator>;
}

impl ContainerConfig {
    pub fn map_uid(&self, internal_uid: u32) -> Option<u32> {
        _map_id(&self.uid_map, internal_uid)
    }
    pub fn map_gid(&self, internal_gid: u32) -> Option<u32> {
        _map_id(&self.gid_map, internal_gid)
    }
    pub fn validator<'x>() -> Box<Validator + 'x> {
        return Box::new(Structure { members: vec!(
            ("kind".to_string(), Box::new(Scalar {
                default: Some("Daemon".to_string()),
                .. Default::default() }) as Box<Validator>),
            ("volumes".to_string(), Box::new(Mapping {
                key_element: Box::new(Scalar {
                    .. Default::default() }) as Box<Validator>,
                value_element: volume_validator(),
            .. Default::default() }) as Box<Validator>),
            ("user_id".to_string(), Box::new(Numeric {
                default: None,
                .. Default::default()}) as Box<Validator>),
            ("group_id".to_string(), Box::new(Numeric {
                default: Some(0),
                .. Default::default()}) as Box<Validator>),
            ("memory_limit".to_string(), Box::new(Numeric {
                default: Some(0x7fffffffffffffffi64),
                .. Default::default()}) as Box<Validator>),
            ("fileno_limit".to_string(), Box::new(Numeric {
                default: Some(1024),
                .. Default::default()}) as Box<Validator>),
            ("cpu_shares".to_string(), Box::new(Numeric {
                default: Some(1024),
                .. Default::default()}) as Box<Validator>),
            ("restart_timeout".to_string(), Box::new(Numeric {
                min: Some(0),
                max: Some(86400),
                default: Some(1),
                .. Default::default()}) as Box<Validator>),
            ("executable".to_string(), Box::new(Scalar {
                .. Default::default() }) as Box<Validator>),
            ("arguments".to_string(), Box::new(Sequence {
                element: Box::new(Scalar {
                    .. Default::default() }) as Box<Validator>,
                .. Default::default() }) as Box<Validator>),
            ("environ".to_string(), Box::new(Mapping {
                key_element: Box::new(Scalar {
                    .. Default::default() }) as Box<Validator>,
                value_element: Box::new(Scalar {
                    .. Default::default() }) as Box<Validator>,
            .. Default::default() }) as Box<Validator>),
            ("workdir".to_string(), Box::new(Scalar {
                default: Some("/".to_string()),
                .. Default::default()}) as Box<Validator>),
            ("resolv_conf".to_string(), Box::new(Structure { members: vec!(
                ("copy_from_host".to_string(), Box::new(Scalar {
                    default: Some("true".to_string()),
                    .. Default::default() }) as Box<Validator>),
                ), .. Default::default()}) as Box<Validator>),
            ("hosts_file".to_string(), Box::new(Structure { members: vec!(
                ("localhost".to_string(), Box::new(Scalar {
                    default: Some("true".to_string()),
                    .. Default::default() }) as Box<Validator>),
                ("public_hostname".to_string(), Box::new(Scalar {
                    default: Some("true".to_string()),
                    .. Default::default() }) as Box<Validator>),
                ), .. Default::default()}) as Box<Validator>),
            ("uid_map".to_string(), mapping_validator()),
            ("gid_map".to_string(), mapping_validator()),
            ("stdout_stderr_file".to_string(), Box::new(Scalar {
                optional: true,
                default: None,
                .. Default::default() }) as Box<Validator>),
            ("restart_process_only".to_string(), Box::new(Scalar {
                default: Some("false".to_string()),
                .. Default::default() }) as Box<Validator>),
        ), .. Default::default() }) as Box<Validator>;
    }
}

pub fn volume_validator<'a>() -> Box<Validator + 'a> {
    return Box::new(Enum { options: vec!(
        ("Persistent".to_string(),  Box::new(Structure { members: vec!(
            ("path".to_string(),  Box::new(Scalar {
                default: Some("/".to_string()),
                .. Default::default()}) as Box<Validator>),
            ("mkdir".to_string(),  Box::new(Scalar {
                default: Some("false".to_string()),
                .. Default::default()}) as Box<Validator>),
            ("mode".to_string(),  Box::new(Numeric {
                min: Some(0),
                max: Some(0o700),
                default: Some(0o766),
                .. Default::default()}) as Box<Validator>),
            ("user".to_string(),  Box::new(Numeric {
                default: Some(0),
                .. Default::default()}) as Box<Validator>),
            ("group".to_string(),  Box::new(Numeric {
                default: Some(0),
                .. Default::default()}) as Box<Validator>),
            ),.. Default::default()}) as Box<Validator>),
        ("Readonly".to_string(), Box::new(Scalar {
            .. Default::default()}) as Box<Validator>),
        ("Tmpfs".to_string(),  Box::new(Structure { members: vec!(
            ("size".to_string(),  Box::new(Numeric {
                min: Some(0),
                default: Some(100*1024*1024),
                .. Default::default()}) as Box<Validator>),
            ("mode".to_string(),  Box::new(Numeric {
                min: Some(0),
                max: Some(0o1777),
                default: Some(0o766),
                .. Default::default()}) as Box<Validator>),
            ),.. Default::default()}) as Box<Validator>),
        ("Statedir".to_string(),  Box::new(Structure { members: vec!(
            ("path".to_string(),  Box::new(Scalar {
                default: Some("/".to_string()),
                .. Default::default()}) as Box<Validator>),
            ("mode".to_string(),  Box::new(Numeric {
                min: Some(0),
                max: Some(0o700),
                default: Some(0o766),
                .. Default::default()}) as Box<Validator>),
            ("user".to_string(),  Box::new(Numeric {
                default: Some(0),
                .. Default::default()}) as Box<Validator>),
            ("group".to_string(),  Box::new(Numeric {
                default: Some(0),
                .. Default::default()}) as Box<Validator>),
            ),.. Default::default()}) as Box<Validator>),
        ), .. Default::default()}) as Box<Validator>;
}

fn _map_id(map: &Vec<IdMap>, /*internal*/id: u32) -> Option<u32> {
    if map.len() == 0 {
        return Some(id);
    }
    for rng in map.iter() {
        if id >= rng.inside && id <= rng.inside + rng.count {
            return Some(rng.outside + (id - rng.inside));
        }
    }
    None
}
<commit_msg>Fix default permissions for volumes (need execute permissions)<commit_after>use std::path::PathBuf;
use std::default::Default;
use std::collections::BTreeMap;

use quire::validate::{Validator, Structure, Sequence, Scalar, Numeric, Enum};
use quire::validate::{Mapping};
use self::Volume::*;

#[derive(RustcDecodable, RustcEncodable, Clone, PartialEq, Eq)]
pub struct TmpfsInfo {
    pub size: usize,
    pub mode: u32,
}

#[derive(RustcDecodable, RustcEncodable, Clone, PartialEq, Eq)]
pub struct PersistentInfo {
    pub path: PathBuf,
    pub mkdir: bool,
    pub mode: u32,
    pub user: u32,
    pub group: u32,
}

#[derive(RustcDecodable, RustcEncodable, Clone, PartialEq, Eq)]
pub struct StatedirInfo {
    pub path: PathBuf,
    pub mode: u32,
    pub user: u32,
    pub group: u32,
}

#[derive(RustcDecodable, RustcEncodable, Clone, PartialEq, Eq)]
pub enum Volume {
    Readonly(PathBuf),
    Persistent(PersistentInfo),
    Tmpfs(TmpfsInfo),
    Statedir(StatedirInfo),
}

#[derive(RustcDecodable, RustcEncodable, Debug, PartialEq, Eq)]
pub enum ContainerKind {
    Daemon,
    Command,
}

#[derive(RustcDecodable, RustcEncodable)]
pub struct ResolvConf {
    pub copy_from_host: bool,
}

#[derive(RustcDecodable, RustcEncodable)]
pub struct HostsFile {
    pub localhost: bool,
    pub public_hostname: bool,
}

#[derive(RustcDecodable, RustcEncodable, Clone, Copy, Debug)]
pub struct IdMap {
    pub inside: u32,
    pub outside: u32,
    pub count: u32,
}

#[derive(RustcDecodable, RustcEncodable)]
pub struct ContainerConfig {
    pub kind: ContainerKind,
    pub volumes: BTreeMap<String, Volume>,
    pub user_id: u32,
    pub group_id: u32,
    pub restart_timeout: f32,
    pub memory_limit: u64,
    pub fileno_limit: u64,
    pub cpu_shares: usize,
    pub executable: String,
    pub arguments: Vec<String>,
    pub environ: BTreeMap<String, String>,
    pub workdir: PathBuf,
    pub resolv_conf: ResolvConf,
    pub hosts_file: HostsFile,
    pub uid_map: Vec<IdMap>,
    pub gid_map: Vec<IdMap>,
    pub stdout_stderr_file: Option<PathBuf>,
    pub restart_process_only: bool,
}

fn mapping_validator<'x>() -> Box<Validator + 'x> {
    return Box::new(Sequence {
        element: Box::new(Structure { members: vec!(
            ("inside".to_string(), Box::new(Numeric {
                default: None,
                .. Default::default() }) as Box<Validator>),
            ("outside".to_string(), Box::new(Numeric {
                default: None,
                .. Default::default() }) as Box<Validator>),
            ("count".to_string(), Box::new(Numeric {
                default: None,
                .. Default::default() }) as Box<Validator>),
            ), .. Default::default() }) as Box<Validator>,
        .. Default::default() }) as Box<Validator>;
}

impl ContainerConfig {
    pub fn map_uid(&self, internal_uid: u32) -> Option<u32> {
        _map_id(&self.uid_map, internal_uid)
    }
    pub fn map_gid(&self, internal_gid: u32) -> Option<u32> {
        _map_id(&self.gid_map, internal_gid)
    }
    pub fn validator<'x>() -> Box<Validator + 'x> {
        return Box::new(Structure { members: vec!(
            ("kind".to_string(), Box::new(Scalar {
                default: Some("Daemon".to_string()),
                .. Default::default() }) as Box<Validator>),
            ("volumes".to_string(), Box::new(Mapping {
                key_element: Box::new(Scalar {
                    .. Default::default() }) as Box<Validator>,
                value_element: volume_validator(),
            .. Default::default() }) as Box<Validator>),
            ("user_id".to_string(), Box::new(Numeric {
                default: None,
                .. Default::default()}) as Box<Validator>),
            ("group_id".to_string(), Box::new(Numeric {
                default: Some(0),
                .. Default::default()}) as Box<Validator>),
            ("memory_limit".to_string(), Box::new(Numeric {
                default: Some(0x7fffffffffffffffi64),
                .. Default::default()}) as Box<Validator>),
            ("fileno_limit".to_string(), Box::new(Numeric {
                default: Some(1024),
                .. Default::default()}) as Box<Validator>),
            ("cpu_shares".to_string(), Box::new(Numeric {
                default: Some(1024),
                .. Default::default()}) as Box<Validator>),
            ("restart_timeout".to_string(), Box::new(Numeric {
                min: Some(0),
                max: Some(86400),
                default: Some(1),
                .. Default::default()}) as Box<Validator>),
            ("executable".to_string(), Box::new(Scalar {
                .. Default::default() }) as Box<Validator>),
            ("arguments".to_string(), Box::new(Sequence {
                element: Box::new(Scalar {
                    .. Default::default() }) as Box<Validator>,
                .. Default::default() }) as Box<Validator>),
            ("environ".to_string(), Box::new(Mapping {
                key_element: Box::new(Scalar {
                    .. Default::default() }) as Box<Validator>,
                value_element: Box::new(Scalar {
                    .. Default::default() }) as Box<Validator>,
            .. Default::default() }) as Box<Validator>),
            ("workdir".to_string(), Box::new(Scalar {
                default: Some("/".to_string()),
                .. Default::default()}) as Box<Validator>),
            ("resolv_conf".to_string(), Box::new(Structure { members: vec!(
                ("copy_from_host".to_string(), Box::new(Scalar {
                    default: Some("true".to_string()),
                    .. Default::default() }) as Box<Validator>),
                ), .. Default::default()}) as Box<Validator>),
            ("hosts_file".to_string(), Box::new(Structure { members: vec!(
                ("localhost".to_string(), Box::new(Scalar {
                    default: Some("true".to_string()),
                    .. Default::default() }) as Box<Validator>),
                ("public_hostname".to_string(), Box::new(Scalar {
                    default: Some("true".to_string()),
                    .. Default::default() }) as Box<Validator>),
                ), .. Default::default()}) as Box<Validator>),
            ("uid_map".to_string(), mapping_validator()),
            ("gid_map".to_string(), mapping_validator()),
            ("stdout_stderr_file".to_string(), Box::new(Scalar {
                optional: true,
                default: None,
                .. Default::default() }) as Box<Validator>),
            ("restart_process_only".to_string(), Box::new(Scalar {
                default: Some("false".to_string()),
                .. Default::default() }) as Box<Validator>),
        ), .. Default::default() }) as Box<Validator>;
    }
}

pub fn volume_validator<'a>() -> Box<Validator + 'a> {
    return Box::new(Enum { options: vec!(
        ("Persistent".to_string(),  Box::new(Structure { members: vec!(
            ("path".to_string(),  Box::new(Scalar {
                default: Some("/".to_string()),
                .. Default::default()}) as Box<Validator>),
            ("mkdir".to_string(),  Box::new(Scalar {
                default: Some("false".to_string()),
                .. Default::default()}) as Box<Validator>),
            ("mode".to_string(),  Box::new(Numeric {
                min: Some(0),
                max: Some(0o1777),
                default: Some(0o777),
                .. Default::default()}) as Box<Validator>),
            ("user".to_string(),  Box::new(Numeric {
                default: Some(0),
                .. Default::default()}) as Box<Validator>),
            ("group".to_string(),  Box::new(Numeric {
                default: Some(0),
                .. Default::default()}) as Box<Validator>),
            ),.. Default::default()}) as Box<Validator>),
        ("Readonly".to_string(), Box::new(Scalar {
            .. Default::default()}) as Box<Validator>),
        ("Tmpfs".to_string(),  Box::new(Structure { members: vec!(
            ("size".to_string(),  Box::new(Numeric {
                min: Some(0),
                default: Some(100*1024*1024),
                .. Default::default()}) as Box<Validator>),
            ("mode".to_string(),  Box::new(Numeric {
                min: Some(0),
                max: Some(0o1777),
                default: Some(0o777),
                .. Default::default()}) as Box<Validator>),
            ),.. Default::default()}) as Box<Validator>),
        ("Statedir".to_string(),  Box::new(Structure { members: vec!(
            ("path".to_string(),  Box::new(Scalar {
                default: Some("/".to_string()),
                .. Default::default()}) as Box<Validator>),
            ("mode".to_string(),  Box::new(Numeric {
                min: Some(0),
                max: Some(0o1777),
                default: Some(0o777),
                .. Default::default()}) as Box<Validator>),
            ("user".to_string(),  Box::new(Numeric {
                default: Some(0),
                .. Default::default()}) as Box<Validator>),
            ("group".to_string(),  Box::new(Numeric {
                default: Some(0),
                .. Default::default()}) as Box<Validator>),
            ),.. Default::default()}) as Box<Validator>),
        ), .. Default::default()}) as Box<Validator>;
}

fn _map_id(map: &Vec<IdMap>, /*internal*/id: u32) -> Option<u32> {
    if map.len() == 0 {
        return Some(id);
    }
    for rng in map.iter() {
        if id >= rng.inside && id <= rng.inside + rng.count {
            return Some(rng.outside + (id - rng.inside));
        }
    }
    None
}
<|endoftext|>


<commit_before>use cpu::common_defs::OpcodeExecInfo;
use cpu::CpuState;
use cpu::DecodeRegister;
use cpu::common_defs::address_mode::AddressMode;
use cpu::common_defs::opcode_class::OpcodeClass;
use memory::Memory;
use std::collections::HashMap;

macro_rules! set_zs {
    ($cpu_state:expr,$val:expr) => ($cpu_state.Z = $val == 0;$cpu_state.S = $val >= 0x80 );
}

macro_rules! stack_push8 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write8(0x0100 | $cpu_state.sp,$val)); $cpu_state.sp-=1;);
}
macro_rules! stack_push16 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write16(0x0100 as u16 | $cpu_state.sp as u16 -1,$val)); $cpu_state.sp-=2;);
}
macro_rules! stack_pull8 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.pc+=1;
            $mem.read8(0x0100 | $cpu_state.sp;)
        }
    }
}
macro_rules! stack_pull16 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.pc+=2;
            $mem.read16(0x0100 as u16 | $cpu_state.sp as u16 -1)
        }
    }
}

#[derive(Debug)]
pub enum ExecutionError {
  MemoryError(::memory::MemoryError),
  UnexpectedOpcode(String),
  UnexpectedAddressMode(String)
}
impl From<::memory::MemoryError> for ExecutionError {
    fn from(err: ::memory::MemoryError) -> ExecutionError {
        ExecutionError::MemoryError(err)
    }
}

pub struct CpuExecutor {
    op_table: HashMap<u8,OpcodeExecInfo>,
}

impl CpuExecutor {

	/// Construct a new CpuExecutor.
    pub fn new(opcodes: Vec<OpcodeExecInfo> ) -> CpuExecutor {
        // Construct the HashMap and copy it over to CpuExecutor's constructor.
        let mut op_table_temp: HashMap<u8, OpcodeExecInfo> = HashMap::new();
        for oc in opcodes {
            if op_table_temp.contains_key(&oc.opcode) {
                println!("Encountered duplicate opcode: {:#x}", oc.opcode);
                println!("Replacing previous opcode, which was {{ opcode: {:#x}, len: {}, cycles: {}, page_cycles: {} }}", oc.opcode, oc.len, oc.cycles, oc.page_cycles);
            }
            op_table_temp.insert(oc.opcode, oc);
        }

        CpuExecutor {
            op_table: op_table_temp,
        }
    }

	/// Reset the given cpu_state.
    pub fn power_on(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        self.reset(cpu_state,mem);
    }
	/// Reset the given cpu_state.
    pub fn reset(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        cpu_state.pc = mem.read16(0xFFFC).unwrap();
        cpu_state.sp = 0xFD;
        cpu_state.pack_flags(0x24);
    }

    /// Execute a single instruction in the given memory and cpu_state context.
    pub fn step(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) {
        self.fetch_and_decode(cpu_state,mem);
        self.execute(cpu_state,mem);
    }

	/// Fetch the next instruction and perform address resolution.
    pub fn fetch_and_decode(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) -> Result<(),ExecutionError> {
        cpu_state.instruction_register = mem.read8(cpu_state.pc).unwrap();
        cpu_state.decode_register = try!(self.decode(cpu_state,mem));
        cpu_state.pc = cpu_state.pc+1;
        Ok(())
    }

    /// Perform address resolution, returning the info in a DecodeRegister.
    fn decode(self: &CpuExecutor, cpu_state: &CpuState, mem: &Memory) -> Result<DecodeRegister,ExecutionError> {
        let mut dr = DecodeRegister {
            info : self.op_table[&cpu_state.instruction_register].clone(),
            ..Default::default()
        };
        match dr.info.address_mode {
            // no explicit addresses for the following modes
            AddressMode::Accumulator  => { },
            AddressMode::Implied      => { },

            // explicit addresses from here on out
            AddressMode::Absolute => {
                dr.addr_final  = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteX       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.x as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteY       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Immediate       => {
                dr.addr_final  = Some(cpu_state.pc+1);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Indirect        => {
                dr.addr_intermediate = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
            },
            AddressMode::IndexedIndirect => {
                dr.addr_intermediate = Some( (mem.read8(cpu_state.pc+1).unwrap() + cpu_state.x) as u16);
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::IndirectIndexed => {
                dr.addr_init         = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.addr_intermediate = Some(mem.read16(dr.addr_init.unwrap()).unwrap());
                dr.addr_final        = Some( ((dr.addr_intermediate.unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Relative        => {
                dr.addr_final  = Some(cpu_state.pc+1);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPage        => {
                dr.addr_final  = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageX       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.x as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageY       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.y as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            _ => { return Err(ExecutionError::UnexpectedAddressMode(format!("unrecognized addressing mode '{:?}' while decoding instruction_register!",dr.info.address_mode))); }
        }

        return Ok(dr);
    }

    /// Perform the current instruction, returning the CpuState after execution.
    pub fn execute(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) -> Result<(),ExecutionError> {
    	// Figure out which opcode is being executed.
    	match cpu_state.decode_register.info.opcode_class {
    		OpcodeClass::ASL => {
				cpu_state.a = cpu_state.a << 1;
				cpu_state.C = (128 & cpu_state.a) > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LDA => {
    			cpu_state.a = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.a);

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LDX => {
    			cpu_state.x = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.x);

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LDY => {
    			cpu_state.y = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.y);

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LSR => {
    			cpu_state.a = cpu_state.a >> 1;
    			cpu_state.C = 1 & cpu_state.a > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::JMP => {
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::JSR => {
                stack_push16!(cpu_state,mem,cpu_state.pc-1);
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::NOP => {},
    		OpcodeClass::SEC => {
                cpu_state.C = true;
    		},
    		OpcodeClass::STX => {
                try!(mem.write8(cpu_state.decode_register.addr_final.unwrap(),cpu_state.decode_register.value_final.unwrap()));
                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},

			_ => { return Err(ExecutionError::UnexpectedOpcode(format!("Unrecognised opcode class: {:?}", cpu_state.decode_register.info.opcode_class)));}

    	}
        Ok(())
    }
}

<commit_msg>step function now properly returns errors if necessary<commit_after>use cpu::common_defs::OpcodeExecInfo;
use cpu::CpuState;
use cpu::DecodeRegister;
use cpu::common_defs::address_mode::AddressMode;
use cpu::common_defs::opcode_class::OpcodeClass;
use memory::Memory;
use std::collections::HashMap;

macro_rules! set_zs {
    ($cpu_state:expr,$val:expr) => ($cpu_state.Z = $val == 0;$cpu_state.S = $val >= 0x80 );
}

macro_rules! stack_push8 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write8(0x0100 | $cpu_state.sp,$val)); $cpu_state.sp-=1;);
}
macro_rules! stack_push16 {
    ($cpu_state:expr,$mem:expr,$val:expr) => (try!($mem.write16(0x0100 as u16 | $cpu_state.sp as u16 -1,$val)); $cpu_state.sp-=2;);
}
macro_rules! stack_pull8 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.pc+=1;
            $mem.read8(0x0100 | $cpu_state.sp;)
        }
    }
}
macro_rules! stack_pull16 {
    ($cpu_state:expr,$mem:expr) => {
        {
            $cpu_state.pc+=2;
            $mem.read16(0x0100 as u16 | $cpu_state.sp as u16 -1)
        }
    }
}

#[derive(Debug)]
pub enum ExecutionError {
  MemoryError(::memory::MemoryError),
  UnexpectedOpcode(String),
  UnexpectedAddressMode(String)
}
impl From<::memory::MemoryError> for ExecutionError {
    fn from(err: ::memory::MemoryError) -> ExecutionError {
        ExecutionError::MemoryError(err)
    }
}

pub struct CpuExecutor {
    op_table: HashMap<u8,OpcodeExecInfo>,
}

impl CpuExecutor {

	/// Construct a new CpuExecutor.
    pub fn new(opcodes: Vec<OpcodeExecInfo> ) -> CpuExecutor {
        // Construct the HashMap and copy it over to CpuExecutor's constructor.
        let mut op_table_temp: HashMap<u8, OpcodeExecInfo> = HashMap::new();
        for oc in opcodes {
            if op_table_temp.contains_key(&oc.opcode) {
                println!("Encountered duplicate opcode: {:#x}", oc.opcode);
                println!("Replacing previous opcode, which was {{ opcode: {:#x}, len: {}, cycles: {}, page_cycles: {} }}", oc.opcode, oc.len, oc.cycles, oc.page_cycles);
            }
            op_table_temp.insert(oc.opcode, oc);
        }

        CpuExecutor {
            op_table: op_table_temp,
        }
    }

	/// Reset the given cpu_state.
    pub fn power_on(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        self.reset(cpu_state,mem);
    }
	/// Reset the given cpu_state.
    pub fn reset(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) {
        cpu_state.pc = mem.read16(0xFFFC).unwrap();
        cpu_state.sp = 0xFD;
        cpu_state.pack_flags(0x24);
    }

    /// Execute a single instruction in the given memory and cpu_state context.
    pub fn step(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) -> Result<(),ExecutionError> {
        try!(self.fetch_and_decode(cpu_state,mem));
        try!(self.execute(cpu_state,mem));
        Ok(())
    }

	/// Fetch the next instruction and perform address resolution.
    pub fn fetch_and_decode(self: &CpuExecutor, cpu_state: &mut CpuState,mem:&mut Memory) -> Result<(),ExecutionError> {
        cpu_state.instruction_register = mem.read8(cpu_state.pc).unwrap();
        cpu_state.decode_register = try!(self.decode(cpu_state,mem));
        cpu_state.pc = cpu_state.pc+1;
        Ok(())
    }

    /// Perform address resolution, returning the info in a DecodeRegister.
    fn decode(self: &CpuExecutor, cpu_state: &CpuState, mem: &Memory) -> Result<DecodeRegister,ExecutionError> {
        let mut dr = DecodeRegister {
            info : self.op_table[&cpu_state.instruction_register].clone(),
            ..Default::default()
        };
        match dr.info.address_mode {
            // no explicit addresses for the following modes
            AddressMode::Accumulator  => { },
            AddressMode::Implied      => { },

            // explicit addresses from here on out
            AddressMode::Absolute => {
                dr.addr_final  = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteX       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.x as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::AbsoluteY       => {
                dr.addr_final  = Some( ((mem.read16(cpu_state.pc+1).unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Immediate       => {
                dr.addr_final  = Some(cpu_state.pc+1);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Indirect        => {
                dr.addr_intermediate = Some(mem.read16(cpu_state.pc+1).unwrap());
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
            },
            AddressMode::IndexedIndirect => {
                dr.addr_intermediate = Some( (mem.read8(cpu_state.pc+1).unwrap() + cpu_state.x) as u16);
                dr.addr_final        = Some(mem.read16(dr.addr_intermediate.unwrap()).unwrap());
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::IndirectIndexed => {
                dr.addr_init         = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.addr_intermediate = Some(mem.read16(dr.addr_init.unwrap()).unwrap());
                dr.addr_final        = Some( ((dr.addr_intermediate.unwrap() as u32 + cpu_state.y as u32) % 65536) as u16);
                dr.value_final       = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::Relative        => {
                dr.addr_final  = Some(cpu_state.pc+1);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPage        => {
                dr.addr_final  = Some(mem.read8(cpu_state.pc+1).unwrap() as u16);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageX       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.x as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            AddressMode::ZeroPageY       => {
                dr.addr_final  = Some( (mem.read8(cpu_state.pc+1).unwrap() as u16 + cpu_state.y as u16) % 256);
                dr.value_final = Some(mem.read8(dr.addr_final.unwrap()).unwrap());
            },
            _ => { return Err(ExecutionError::UnexpectedAddressMode(format!("unrecognized addressing mode '{:?}' while decoding instruction_register!",dr.info.address_mode))); }
        }

        return Ok(dr);
    }

    /// Perform the current instruction, returning the CpuState after execution.
    pub fn execute(self: &CpuExecutor, cpu_state: &mut CpuState, mem:&mut Memory) -> Result<(),ExecutionError> {
    	// Figure out which opcode is being executed.
    	match cpu_state.decode_register.info.opcode_class {
    		OpcodeClass::ASL => {
				cpu_state.a = cpu_state.a << 1;
				cpu_state.C = (128 & cpu_state.a) > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LDA => {
    			cpu_state.a = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.a);

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LDX => {
    			cpu_state.x = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.x);

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LDY => {
    			cpu_state.y = cpu_state.decode_register.value_final.unwrap();
                set_zs!(cpu_state,cpu_state.y);

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::LSR => {
    			cpu_state.a = cpu_state.a >> 1;
    			cpu_state.C = 1 & cpu_state.a > 0;
    			if cpu_state.a == 0 { cpu_state.Z = true; }

                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},
    		OpcodeClass::JMP => {
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::JSR => {
                stack_push16!(cpu_state,mem,cpu_state.pc-1);
    			cpu_state.pc = cpu_state.decode_register.addr_final.unwrap();
    		},
    		OpcodeClass::NOP => {},
    		OpcodeClass::SEC => {
                cpu_state.C = true;
    		},
    		OpcodeClass::STX => {
                try!(mem.write8(cpu_state.decode_register.addr_final.unwrap(),cpu_state.decode_register.value_final.unwrap()));
                cpu_state.pc = cpu_state.pc + (cpu_state.decode_register.info.len as u16 -1);
    		},

			_ => { return Err(ExecutionError::UnexpectedOpcode(format!("Unrecognised opcode class: {:?}", cpu_state.decode_register.info.opcode_class)));}

    	}
        Ok(())
    }
}

<|endoftext|>


<commit_before>extern {
}

extern crate gb;

use std::ptr;
use std::mem;
use std::ffi::{
    CString,
};
use std::os::raw::{
    c_char,
    c_void,
};
use std::slice;

use gb::{
    Emulator,
    Interrupt,
    Hardware,
    Key,
};

fn store_frame(screen: &gb::ScreenBuffer, data: &mut [u8]) {
    for i in 0..gb::SCREEN_X {
        for j in 0..gb::SCREEN_Y {
            data[i * gb::SCREEN_Y + j] = screen[j][i] as u8;
        }
    }
}

static mut EMULATOR: Option<Emulator> = None;
static mut SCREEN: *mut u8 = 0 as *mut u8;
static mut SOUND: *mut u8 = 0 as *mut u8;
static mut GAMEPAD: &mut [u8] = &mut [];

#[derive(Debug)]
struct GamepadStatus {
    a: bool,
    b: bool,
    start: bool,
    select: bool,
    up: bool,
    down: bool,
    left: bool,
    right: bool,
}

#[no_mangle]
pub extern "C" fn alloc(size: usize) -> *mut c_void {
    let mut buf = Vec::with_capacity(size);
    let ptr = buf.as_mut_ptr();
    mem::forget(buf);
    return ptr as *mut c_void;
}

#[no_mangle]
pub extern "C" fn dealloc(ptr: *mut c_void, cap: usize) {
    unsafe  {
        let _buf = Vec::from_raw_parts(ptr, 0, cap);
    }
}

#[no_mangle]
pub unsafe extern "C" fn main_loop() {
	if EMULATOR.is_none() || GAMEPAD.len() == 0 || SOUND == ptr::null_mut()
        || SCREEN == ptr::null_mut() {
		return;
	}

    let sound = slice::from_raw_parts_mut(SOUND as *mut i16, 1470);
    let screen = slice::from_raw_parts_mut(SCREEN as *mut u8,
                                           gb::SCREEN_X * gb::SCREEN_Y * 3);
	main_loop_internal(EMULATOR.as_mut().unwrap(),
        screen, sound, &GamepadStatus::from_raw(GAMEPAD));
}

impl GamepadStatus {
    fn from_raw(data: &[u8]) -> GamepadStatus {
        GamepadStatus {
            a: data[0] == 1,
            b: data[1] == 1,
            start: data[2] == 1,
            select: data[3] == 1,
            up: data[4] == 1,
            down: data[5] == 1,
            left: data[6] == 1,
            right: data[7] == 1,
        }
    }
}

fn update_button(emulator: &mut Emulator, button: Key, pressed: bool) {
    if pressed {
        emulator.cpu.key_down(button);
    } else {
        emulator.cpu.key_up(button);
    }
}

fn main_loop_internal(emulator: &mut Emulator, screen: &mut [u8],
                      sound: &mut [i16], gamepad: &GamepadStatus) {
    update_button(emulator, Key::A, gamepad.a);
    update_button(emulator, Key::B, gamepad.b);
    update_button(emulator, Key::Select, gamepad.select);
    update_button(emulator, Key::Start, gamepad.start);
    update_button(emulator, Key::Up, gamepad.up);
    update_button(emulator, Key::Down, gamepad.down);
    update_button(emulator, Key::Left, gamepad.left);
    update_button(emulator, Key::Right, gamepad.right);
    emulator.cpu.interrupt(Interrupt::Joypad);

	loop {
		emulator.cpu.next_instruction();
		if emulator.cpu.handler_holder.should_refresh() {
			break;
		}
	}

    store_frame(emulator.cpu.handler_holder.get_screen_buffer(),
                screen);

    emulator.generate_sound_into(sound);
}

#[no_mangle]
pub fn init(data: *mut u8, data_size: isize, screen_data: *mut u8,
                   sound_data: *mut u8, gamepad_data: *mut u8) -> *mut c_char {
    unsafe {
        let bytes = slice::from_raw_parts(data, data_size as usize);
		EMULATOR = Some(Emulator::from_data(&bytes, 44100.00).unwrap());
        SCREEN = screen_data;
        SOUND = sound_data;
        GAMEPAD = slice::from_raw_parts_mut(gamepad_data, 8);
	}

    CString::new("OK")
        .unwrap()
        .into_raw()
}

fn main() {
}

<commit_msg>Fix whitespace in wasm.<commit_after>extern {
}

extern crate gb;

use std::ptr;
use std::mem;
use std::ffi::{
    CString,
};
use std::os::raw::{
    c_char,
    c_void,
};
use std::slice;

use gb::{
    Emulator,
    Interrupt,
    Hardware,
    Key,
};

fn store_frame(screen: &gb::ScreenBuffer, data: &mut [u8]) {
    for i in 0..gb::SCREEN_X {
        for j in 0..gb::SCREEN_Y {
            data[i * gb::SCREEN_Y + j] = screen[j][i] as u8;
        }
    }
}

static mut EMULATOR: Option<Emulator> = None;
static mut SCREEN: *mut u8 = 0 as *mut u8;
static mut SOUND: *mut u8 = 0 as *mut u8;
static mut GAMEPAD: &mut [u8] = &mut [];

#[derive(Debug)]
struct GamepadStatus {
    a: bool,
    b: bool,
    start: bool,
    select: bool,
    up: bool,
    down: bool,
    left: bool,
    right: bool,
}

#[no_mangle]
pub extern "C" fn alloc(size: usize) -> *mut c_void {
    let mut buf = Vec::with_capacity(size);
    let ptr = buf.as_mut_ptr();
    mem::forget(buf);
    return ptr as *mut c_void;
}

#[no_mangle]
pub extern "C" fn dealloc(ptr: *mut c_void, cap: usize) {
    unsafe  {
        let _buf = Vec::from_raw_parts(ptr, 0, cap);
    }
}

#[no_mangle]
pub unsafe extern "C" fn main_loop() {
    if EMULATOR.is_none() || GAMEPAD.len() == 0 || SOUND == ptr::null_mut()
        || SCREEN == ptr::null_mut() {
        return;
    }

    let sound = slice::from_raw_parts_mut(SOUND as *mut i16, 1470);
    let screen = slice::from_raw_parts_mut(SCREEN as *mut u8,
                                           gb::SCREEN_X * gb::SCREEN_Y * 3);
    main_loop_internal(EMULATOR.as_mut().unwrap(),
        screen, sound, &GamepadStatus::from_raw(GAMEPAD));
}

impl GamepadStatus {
    fn from_raw(data: &[u8]) -> GamepadStatus {
        GamepadStatus {
            a: data[0] == 1,
            b: data[1] == 1,
            start: data[2] == 1,
            select: data[3] == 1,
            up: data[4] == 1,
            down: data[5] == 1,
            left: data[6] == 1,
            right: data[7] == 1,
        }
    }
}

fn update_button(emulator: &mut Emulator, button: Key, pressed: bool) {
    if pressed {
        emulator.cpu.key_down(button);
    } else {
        emulator.cpu.key_up(button);
    }
}

fn main_loop_internal(emulator: &mut Emulator, screen: &mut [u8],
                      sound: &mut [i16], gamepad: &GamepadStatus) {
    update_button(emulator, Key::A, gamepad.a);
    update_button(emulator, Key::B, gamepad.b);
    update_button(emulator, Key::Select, gamepad.select);
    update_button(emulator, Key::Start, gamepad.start);
    update_button(emulator, Key::Up, gamepad.up);
    update_button(emulator, Key::Down, gamepad.down);
    update_button(emulator, Key::Left, gamepad.left);
    update_button(emulator, Key::Right, gamepad.right);
    emulator.cpu.interrupt(Interrupt::Joypad);

    loop {
        emulator.cpu.next_instruction();
        if emulator.cpu.handler_holder.should_refresh() {
            break;
        }
    }

    store_frame(emulator.cpu.handler_holder.get_screen_buffer(),
                screen);

    emulator.generate_sound_into(sound);
}

#[no_mangle]
pub fn init(data: *mut u8, data_size: isize, screen_data: *mut u8,
                   sound_data: *mut u8, gamepad_data: *mut u8) -> *mut c_char {
    unsafe {
        let bytes = slice::from_raw_parts(data, data_size as usize);
        EMULATOR = Some(Emulator::from_data(&bytes, 44100.00).unwrap());
        SCREEN = screen_data;
        SOUND = sound_data;
        GAMEPAD = slice::from_raw_parts_mut(gamepad_data, 8);
    }

    CString::new("OK")
        .unwrap()
        .into_raw()
}

fn main() {
}

<|endoftext|>


<commit_before>/*!
`ndarray` compability

This adds support for reading data into [`ndarray`][ndarray] array objects. This primarily
enables support for automatically reshaping the resulting arrays to have the dimensionality of
the image that was asked for.

Data is read into the [`ndarray::ArrayD`][arrayd] type. The following methods from
[`FitsHdu`][fits-hdu] are supported:

* [`read_image`][read-image]
* [`read_region`][read-region]
* [`read_row`][read-row]
* [`read_rows`][read-rows]
* [`read_section`][read-section]

## `read_image`

```rust
# extern crate fitsio;
# #[cfg(feature = "array")]
# extern crate ndarray;
use fitsio::FitsFile;
# #[cfg(feature = "array")]
use ndarray::ArrayD;

# #[cfg(feature = "array")]
# fn main() {
let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
let hdu = f.primary_hdu().unwrap();

let data: ArrayD<u32> = hdu.read_image(&mut f).unwrap();
let dim = data.dim();
assert_eq!(data.ndim(), 2);
assert_eq!(dim[0], 100);
assert_eq!(dim[1], 100);
assert_eq!(data[[20, 5]], 152);
# }
#
# #[cfg(not(feature = "array"))]
# fn main() {}
```

## `read_region`

```rust
# extern crate fitsio;
# #[cfg(feature = "array")]
# extern crate ndarray;
use fitsio::FitsFile;
# #[cfg(feature = "array")]
use ndarray::ArrayD;

# #[cfg(feature = "array")]
# fn main() {
let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
let hdu = f.primary_hdu().unwrap();

let data: ArrayD<u32> = hdu.read_region(&mut f, &[&(70..80), &(20..50)]).unwrap();
let dim = data.dim();
assert_eq!(data.ndim(), 2);
assert_eq!(dim[0], 10);
assert_eq!(dim[1], 30);
assert_eq!(data[[5, 10]], 177);
# }
#
# #[cfg(not(feature = "array"))]
# fn main() {}
```

## `read_row`

```rust
# extern crate fitsio;
# #[cfg(feature = "array")]
# extern crate ndarray;
use fitsio::FitsFile;
# #[cfg(feature = "array")]
use ndarray::ArrayD;

# #[cfg(feature = "array")]
# fn main() {
let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
let hdu = f.primary_hdu().unwrap();

let data: ArrayD<u32> = hdu.read_row(&mut f, 49).unwrap();
assert_eq!(data.ndim(), 1);
assert_eq!(data[20], 156);
# }
#
# #[cfg(not(feature = "array"))]
# fn main() {}
```

## `read_rows`

```rust
# extern crate fitsio;
# #[cfg(feature = "array")]
# extern crate ndarray;
use fitsio::FitsFile;
# #[cfg(feature = "array")]
use ndarray::ArrayD;

# #[cfg(feature = "array")]
# fn main() {
let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
let hdu = f.primary_hdu().unwrap();

let data: ArrayD<u32> = hdu.read_rows(&mut f, 0, 2).unwrap();
let dim = data.dim();
assert_eq!(data.ndim(), 2);
assert_eq!(dim[0], 2);
assert_eq!(dim[1], 100);
assert_eq!(data[[1, 52]], 184);
# }
#
# #[cfg(not(feature = "array"))]
# fn main() {}
```

## `read_section`

```rust
# extern crate fitsio;
# #[cfg(feature = "array")]
# extern crate ndarray;
use fitsio::FitsFile;
use fitsio::errors::Error;
# #[cfg(feature = "array")]
use ndarray::ArrayD;

# #[cfg(feature = "array")]
# fn main() {
let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
let hdu = f.primary_hdu().unwrap();

match hdu.read_section::<ArrayD<u32>>(&mut f, 0, 250) {
    Err(Error::Message(msg)) => {
        assert_eq!(
            msg,
            "must request number of pixels exactly divisible by image width".to_string()
        );
    }
    _ => panic!("invalid result"),
}

let data: ArrayD<u32> = hdu.read_section(&mut f, 0, 200).unwrap();
let dim = data.dim();
assert_eq!(data.ndim(), 2);
assert_eq!(dim[0], 2);
assert_eq!(dim[1], 100);
assert_eq!(data[[0, 10]], 160);
# }
#
# #[cfg(not(feature = "array"))]
# fn main() {}
```

[ndarray]: https://crates.io/crates/ndarray
[arrayd]: https://docs.rs/ndarray/0.11.2/ndarray/type.ArrayD.html
[fits-hdu]: hdu/struct.FitsHdu.html
[read-image]: images/struct.FitsHdu.html#method.read_image
[read-region]: images/struct.FitsHdu.html#method.read_region
[read-row]: images/struct.FitsHdu.html#method.read_row
[read-rows]: images/struct.FitsHdu.html#method.read_rows
[read-section]: images/struct.FitsHdu.html#method.read_section
*/

use crate::errors::Result;
use crate::fitsfile::FitsFile;
use crate::hdu::{FitsHdu, HduInfo};
use crate::images::ReadImage;
use ndarray::{Array, ArrayD};
use std::ops::Range;

impl<T> ReadImage for ArrayD<T>
where
    T: Clone,
    Vec<T>: ReadImage,
{
    fn read_section(fits_file: &mut FitsFile, hdu: &FitsHdu, range: Range<usize>) -> Result<Self> {
        match hdu.info {
            HduInfo::ImageInfo { ref shape, .. } => {
                if shape.len() != 2 {
                    return Err("Only 2D images supported for now".into());
                }

                let width = shape[1];

                if range.start % width != 0 {
                    return Err("range must start on row boundary".into());
                }
                let start_pixel = range.start / width;

                let n_pixels_requested = range.end - range.start;
                if n_pixels_requested % width != 0 {
                    return Err(
                        "must request number of pixels exactly divisible by image width".into(),
                    );
                }

                let n_rows = n_pixels_requested / width;
                ReadImage::read_rows(fits_file, hdu, start_pixel, n_rows)
            }
            HduInfo::TableInfo { .. } => {
                return Err("Cannot read image data from a FITS table".into());
            }
            _ => unreachable!(),
        }
    }

    fn read_rows(
        fits_file: &mut FitsFile,
        hdu: &FitsHdu,
        start_row: usize,
        num_rows: usize,
    ) -> Result<Self> {
        let data: Vec<T> = ReadImage::read_rows(fits_file, hdu, start_row, num_rows)?;
        let arr = Array::from_vec(data);
        let row_length = arr.len() / num_rows;
        Ok(arr.into_shape(vec![num_rows, row_length]).unwrap())
    }

    fn read_row(fits_file: &mut FitsFile, hdu: &FitsHdu, row: usize) -> Result<Self> {
        let data: Vec<T> = ReadImage::read_row(fits_file, hdu, row)?;
        let shape = vec![data.len()];
        Ok(Array::from_shape_vec(shape, data).unwrap())
    }

    fn read_region(
        fits_file: &mut FitsFile,
        hdu: &FitsHdu,
        ranges: &[&Range<usize>],
    ) -> Result<Self> {
        let data: Vec<T> = ReadImage::read_region(fits_file, hdu, ranges)?;
        let shape: Vec<usize> = (0..ranges.len())
            .map(|i| ranges[i].end - ranges[i].start)
            .collect();
        let arr = Array::from_shape_vec(shape, data).unwrap();
        Ok(arr)
    }

    fn read_image(fits_file: &mut FitsFile, hdu: &FitsHdu) -> Result<Self> {
        match hdu.info {
            HduInfo::ImageInfo { ref shape, .. } => {
                let data: Vec<T> = ReadImage::read_image(fits_file, hdu)?;
                let shape: Vec<usize> = (0..2).map(|i| shape[i]).collect();
                let arr = Array::from_shape_vec(shape, data).unwrap();
                Ok(arr)
            }
            _ => unreachable!(),
        }
    }
}

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

    #[test]
    fn test_read_image() {
        let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
        let hdu = f.primary_hdu().unwrap();

        let data: ArrayD<u32> = hdu.read_image(&mut f).unwrap();
        let dim = data.dim();
        assert_eq!(data.ndim(), 2);
        assert_eq!(dim[0], 100);
        assert_eq!(dim[1], 100);
        assert_eq!(data[[20, 5]], 152);
    }

    #[test]
    fn test_read_rows() {
        let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
        let hdu = f.primary_hdu().unwrap();

        let data: ArrayD<u32> = hdu.read_rows(&mut f, 0, 2).unwrap();
        let dim = data.dim();
        assert_eq!(data.ndim(), 2);
        assert_eq!(dim[0], 2);
        assert_eq!(dim[1], 100);
        assert_eq!(data[[1, 52]], 184);
    }

    #[test]
    fn test_read_row() {
        let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
        let hdu = f.primary_hdu().unwrap();

        let data: ArrayD<u32> = hdu.read_row(&mut f, 49).unwrap();
        assert_eq!(data.ndim(), 1);
        assert_eq!(data[20], 156);
    }

    #[test]
    fn test_read_region() {
        let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
        let hdu = f.primary_hdu().unwrap();

        let data: ArrayD<u32> = hdu.read_region(&mut f, &[&(70..80), &(20..50)]).unwrap();
        let dim = data.dim();
        assert_eq!(data.ndim(), 2);
        assert_eq!(dim[0], 10);
        assert_eq!(dim[1], 30);
        assert_eq!(data[[5, 10]], 177);
    }

    #[test]
    fn test_read_section() {
        let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
        let hdu = f.primary_hdu().unwrap();

        match hdu.read_section::<ArrayD<u32>>(&mut f, 0, 250) {
            Err(Error::Message(msg)) => {
                assert_eq!(
                    msg,
                    "must request number of pixels exactly divisible by image width".to_string()
                );
            }
            _ => panic!("invalid result"),
        }

        let data: ArrayD<u32> = hdu.read_section(&mut f, 0, 200).unwrap();
        let dim = data.dim();
        assert_eq!(data.ndim(), 2);
        assert_eq!(dim[0], 2);
        assert_eq!(dim[1], 100);
        assert_eq!(data[[0, 10]], 160);
    }
}
<commit_msg>Fixes deprecation warning for `ndarray` crate<commit_after>/*!
`ndarray` compability

This adds support for reading data into [`ndarray`][ndarray] array objects. This primarily
enables support for automatically reshaping the resulting arrays to have the dimensionality of
the image that was asked for.

Data is read into the [`ndarray::ArrayD`][arrayd] type. The following methods from
[`FitsHdu`][fits-hdu] are supported:

* [`read_image`][read-image]
* [`read_region`][read-region]
* [`read_row`][read-row]
* [`read_rows`][read-rows]
* [`read_section`][read-section]

## `read_image`

```rust
# extern crate fitsio;
# #[cfg(feature = "array")]
# extern crate ndarray;
use fitsio::FitsFile;
# #[cfg(feature = "array")]
use ndarray::ArrayD;

# #[cfg(feature = "array")]
# fn main() {
let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
let hdu = f.primary_hdu().unwrap();

let data: ArrayD<u32> = hdu.read_image(&mut f).unwrap();
let dim = data.dim();
assert_eq!(data.ndim(), 2);
assert_eq!(dim[0], 100);
assert_eq!(dim[1], 100);
assert_eq!(data[[20, 5]], 152);
# }
#
# #[cfg(not(feature = "array"))]
# fn main() {}
```

## `read_region`

```rust
# extern crate fitsio;
# #[cfg(feature = "array")]
# extern crate ndarray;
use fitsio::FitsFile;
# #[cfg(feature = "array")]
use ndarray::ArrayD;

# #[cfg(feature = "array")]
# fn main() {
let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
let hdu = f.primary_hdu().unwrap();

let data: ArrayD<u32> = hdu.read_region(&mut f, &[&(70..80), &(20..50)]).unwrap();
let dim = data.dim();
assert_eq!(data.ndim(), 2);
assert_eq!(dim[0], 10);
assert_eq!(dim[1], 30);
assert_eq!(data[[5, 10]], 177);
# }
#
# #[cfg(not(feature = "array"))]
# fn main() {}
```

## `read_row`

```rust
# extern crate fitsio;
# #[cfg(feature = "array")]
# extern crate ndarray;
use fitsio::FitsFile;
# #[cfg(feature = "array")]
use ndarray::ArrayD;

# #[cfg(feature = "array")]
# fn main() {
let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
let hdu = f.primary_hdu().unwrap();

let data: ArrayD<u32> = hdu.read_row(&mut f, 49).unwrap();
assert_eq!(data.ndim(), 1);
assert_eq!(data[20], 156);
# }
#
# #[cfg(not(feature = "array"))]
# fn main() {}
```

## `read_rows`

```rust
# extern crate fitsio;
# #[cfg(feature = "array")]
# extern crate ndarray;
use fitsio::FitsFile;
# #[cfg(feature = "array")]
use ndarray::ArrayD;

# #[cfg(feature = "array")]
# fn main() {
let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
let hdu = f.primary_hdu().unwrap();

let data: ArrayD<u32> = hdu.read_rows(&mut f, 0, 2).unwrap();
let dim = data.dim();
assert_eq!(data.ndim(), 2);
assert_eq!(dim[0], 2);
assert_eq!(dim[1], 100);
assert_eq!(data[[1, 52]], 184);
# }
#
# #[cfg(not(feature = "array"))]
# fn main() {}
```

## `read_section`

```rust
# extern crate fitsio;
# #[cfg(feature = "array")]
# extern crate ndarray;
use fitsio::FitsFile;
use fitsio::errors::Error;
# #[cfg(feature = "array")]
use ndarray::ArrayD;

# #[cfg(feature = "array")]
# fn main() {
let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
let hdu = f.primary_hdu().unwrap();

match hdu.read_section::<ArrayD<u32>>(&mut f, 0, 250) {
    Err(Error::Message(msg)) => {
        assert_eq!(
            msg,
            "must request number of pixels exactly divisible by image width".to_string()
        );
    }
    _ => panic!("invalid result"),
}

let data: ArrayD<u32> = hdu.read_section(&mut f, 0, 200).unwrap();
let dim = data.dim();
assert_eq!(data.ndim(), 2);
assert_eq!(dim[0], 2);
assert_eq!(dim[1], 100);
assert_eq!(data[[0, 10]], 160);
# }
#
# #[cfg(not(feature = "array"))]
# fn main() {}
```

[ndarray]: https://crates.io/crates/ndarray
[arrayd]: https://docs.rs/ndarray/0.11.2/ndarray/type.ArrayD.html
[fits-hdu]: hdu/struct.FitsHdu.html
[read-image]: images/struct.FitsHdu.html#method.read_image
[read-region]: images/struct.FitsHdu.html#method.read_region
[read-row]: images/struct.FitsHdu.html#method.read_row
[read-rows]: images/struct.FitsHdu.html#method.read_rows
[read-section]: images/struct.FitsHdu.html#method.read_section
*/

use crate::errors::Result;
use crate::fitsfile::FitsFile;
use crate::hdu::{FitsHdu, HduInfo};
use crate::images::ReadImage;
use ndarray::{Array, ArrayD};
use std::ops::Range;

impl<T> ReadImage for ArrayD<T>
where
    T: Clone,
    Vec<T>: ReadImage,
{
    fn read_section(fits_file: &mut FitsFile, hdu: &FitsHdu, range: Range<usize>) -> Result<Self> {
        match hdu.info {
            HduInfo::ImageInfo { ref shape, .. } => {
                if shape.len() != 2 {
                    return Err("Only 2D images supported for now".into());
                }

                let width = shape[1];

                if range.start % width != 0 {
                    return Err("range must start on row boundary".into());
                }
                let start_pixel = range.start / width;

                let n_pixels_requested = range.end - range.start;
                if n_pixels_requested % width != 0 {
                    return Err(
                        "must request number of pixels exactly divisible by image width".into(),
                    );
                }

                let n_rows = n_pixels_requested / width;
                ReadImage::read_rows(fits_file, hdu, start_pixel, n_rows)
            }
            HduInfo::TableInfo { .. } => {
                return Err("Cannot read image data from a FITS table".into());
            }
            _ => unreachable!(),
        }
    }

    fn read_rows(
        fits_file: &mut FitsFile,
        hdu: &FitsHdu,
        start_row: usize,
        num_rows: usize,
    ) -> Result<Self> {
        let data: Vec<T> = ReadImage::read_rows(fits_file, hdu, start_row, num_rows)?;
        let arr = Array::from(data);
        let row_length = arr.len() / num_rows;
        Ok(arr.into_shape(vec![num_rows, row_length]).unwrap())
    }

    fn read_row(fits_file: &mut FitsFile, hdu: &FitsHdu, row: usize) -> Result<Self> {
        let data: Vec<T> = ReadImage::read_row(fits_file, hdu, row)?;
        let shape = vec![data.len()];
        Ok(Array::from_shape_vec(shape, data).unwrap())
    }

    fn read_region(
        fits_file: &mut FitsFile,
        hdu: &FitsHdu,
        ranges: &[&Range<usize>],
    ) -> Result<Self> {
        let data: Vec<T> = ReadImage::read_region(fits_file, hdu, ranges)?;
        let shape: Vec<usize> = (0..ranges.len())
            .map(|i| ranges[i].end - ranges[i].start)
            .collect();
        let arr = Array::from_shape_vec(shape, data).unwrap();
        Ok(arr)
    }

    fn read_image(fits_file: &mut FitsFile, hdu: &FitsHdu) -> Result<Self> {
        match hdu.info {
            HduInfo::ImageInfo { ref shape, .. } => {
                let data: Vec<T> = ReadImage::read_image(fits_file, hdu)?;
                let shape: Vec<usize> = (0..2).map(|i| shape[i]).collect();
                let arr = Array::from_shape_vec(shape, data).unwrap();
                Ok(arr)
            }
            _ => unreachable!(),
        }
    }
}

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

    #[test]
    fn test_read_image() {
        let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
        let hdu = f.primary_hdu().unwrap();

        let data: ArrayD<u32> = hdu.read_image(&mut f).unwrap();
        let dim = data.dim();
        assert_eq!(data.ndim(), 2);
        assert_eq!(dim[0], 100);
        assert_eq!(dim[1], 100);
        assert_eq!(data[[20, 5]], 152);
    }

    #[test]
    fn test_read_rows() {
        let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
        let hdu = f.primary_hdu().unwrap();

        let data: ArrayD<u32> = hdu.read_rows(&mut f, 0, 2).unwrap();
        let dim = data.dim();
        assert_eq!(data.ndim(), 2);
        assert_eq!(dim[0], 2);
        assert_eq!(dim[1], 100);
        assert_eq!(data[[1, 52]], 184);
    }

    #[test]
    fn test_read_row() {
        let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
        let hdu = f.primary_hdu().unwrap();

        let data: ArrayD<u32> = hdu.read_row(&mut f, 49).unwrap();
        assert_eq!(data.ndim(), 1);
        assert_eq!(data[20], 156);
    }

    #[test]
    fn test_read_region() {
        let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
        let hdu = f.primary_hdu().unwrap();

        let data: ArrayD<u32> = hdu.read_region(&mut f, &[&(70..80), &(20..50)]).unwrap();
        let dim = data.dim();
        assert_eq!(data.ndim(), 2);
        assert_eq!(dim[0], 10);
        assert_eq!(dim[1], 30);
        assert_eq!(data[[5, 10]], 177);
    }

    #[test]
    fn test_read_section() {
        let mut f = FitsFile::open("../testdata/full_example.fits").unwrap();
        let hdu = f.primary_hdu().unwrap();

        match hdu.read_section::<ArrayD<u32>>(&mut f, 0, 250) {
            Err(Error::Message(msg)) => {
                assert_eq!(
                    msg,
                    "must request number of pixels exactly divisible by image width".to_string()
                );
            }
            _ => panic!("invalid result"),
        }

        let data: ArrayD<u32> = hdu.read_section(&mut f, 0, 200).unwrap();
        let dim = data.dim();
        assert_eq!(data.ndim(), 2);
        assert_eq!(dim[0], 2);
        assert_eq!(dim[1], 100);
        assert_eq!(data[[0, 10]], 160);
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright 2017-2019 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
 * http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
 * <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
 * option. This file may not be copied, modified, or distributed
 * except according to those terms.
 */

use std::collections::{hash_map::Entry, HashMap};
use std::net::SocketAddr;

use futures::{
    future,
    stream::Fuse,
    sync::{mpsc, oneshot},
    Async, AsyncSink, Future, Poll, Sink, Stream,
};

use tokio_executor::{DefaultExecutor, Executor};

use super::connect::{connect, RespConnection};
use error;
use reconnect::{reconnect, Reconnect};
use resp::{self, FromResp};

#[derive(Debug)]
enum PubsubEvent {
    Subscribe(String, PubsubSink, oneshot::Sender<()>),
    Unsubscribe(String),
}

pub type PubsubStreamInner = mpsc::UnboundedReceiver<resp::RespValue>;
pub type PubsubSink = mpsc::UnboundedSender<resp::RespValue>;

struct PubsubConnectionInner {
    connection: RespConnection,
    out_rx: Fuse<mpsc::UnboundedReceiver<PubsubEvent>>,
    subscriptions: HashMap<String, PubsubSink>,
    pending_subs: HashMap<String, (PubsubSink, oneshot::Sender<()>)>,
    send_pending: Option<resp::RespValue>,
}

impl PubsubConnectionInner {
    fn new(con: RespConnection, out_rx: mpsc::UnboundedReceiver<PubsubEvent>) -> Self {
        PubsubConnectionInner {
            connection: con,
            out_rx: out_rx.fuse(),
            subscriptions: HashMap::new(),
            pending_subs: HashMap::new(),
            send_pending: None,
        }
    }

    /// Returns true = OK, more can be sent, or false = sink is full, needs flushing
    fn do_send(&mut self, msg: resp::RespValue) -> Result<bool, ()> {
        match self
            .connection
            .start_send(msg)
            .map_err(|e| error!("Cannot send subscription request to Redis: {}", e))?
        {
            AsyncSink::Ready => Ok(true),
            AsyncSink::NotReady(msg) => {
                self.send_pending = Some(msg);
                Ok(false)
            }
        }
    }

    fn do_flush(&mut self) -> Result<(), ()> {
        self.connection
            .poll_complete()
            .map(|_| ())
            .map_err(|e| error!("Error polling for completeness: {}", e))
    }

    // Returns true = flushing required.  false = no flushing required
    fn handle_new_subs(&mut self) -> Result<bool, ()> {
        let mut flushing_req = false;
        if let Some(msg) = self.send_pending.take() {
            flushing_req = true;
            if !self.do_send(msg)? {
                return Ok(flushing_req);
            }
        }
        loop {
            match self
                .out_rx
                .poll()
                .map_err(|_| error!("Cannot poll for new subscriptions"))?
            {
                Async::Ready(None) => {
                    return Ok(flushing_req);
                }
                Async::Ready(Some(pubsub_event)) => {
                    let message = match pubsub_event {
                        PubsubEvent::Subscribe(topic, sender, signal) => {
                            self.pending_subs.insert(topic.clone(), (sender, signal));
                            resp_array!["SUBSCRIBE", topic]
                        }
                        PubsubEvent::Unsubscribe(topic) => resp_array!["UNSUBSCRIBE", topic],
                    };
                    flushing_req = true;
                    if !self.do_send(message)? {
                        return Ok(flushing_req);
                    }
                }
                Async::NotReady => {
                    return Ok(flushing_req);
                }
            }
        }
    }

    fn handle_message(&mut self, msg: resp::RespValue) -> Result<bool, ()> {
        let (message_type, topic, msg) = match msg {
            resp::RespValue::Array(mut messages) => match (
                messages.pop(),
                messages.pop(),
                messages.pop(),
                messages.pop(),
            ) {
                (Some(msg), Some(topic), Some(message_type), None) => {
                    match (msg, String::from_resp(topic), message_type) {
                        (msg, Ok(topic), resp::RespValue::BulkString(bytes)) => (bytes, topic, msg),
                        _ => {
                            error!("Incorrect format of PUBSUB message");
                            return Err(());
                        }
                    }
                }
                _ => {
                    error!("Wrong number of parts for a PUBSUB message");
                    return Err(());
                }
            },
            _ => {
                error!("PUBSUB message should be encoded as an array");
                return Err(());
            }
        };

        if message_type == b"subscribe" {
            if let Some((sender, signal)) = self.pending_subs.remove(&topic) {
                self.subscriptions.insert(topic, sender);
                signal
                    .send(())
                    .map_err(|_| error!("Error confirming subscription"))?;
            }
        } else if message_type == b"unsubscribe" {
            if let Entry::Occupied(entry) = self.subscriptions.entry(topic) {
                entry.remove_entry();
            }
            if self.subscriptions.is_empty() {
                return Ok(false);
            }
        } else if message_type == b"message" {
            if let Some(sender) = self.subscriptions.get(&topic) {
                sender.unbounded_send(msg).expect("Cannot send message");
            }
        }

        Ok(true)
    }

    /// Returns true, if there are still valid subscriptions at the end, or false if not, i.e. the whole thing can be dropped.
    fn handle_messages(&mut self) -> Result<bool, ()> {
        loop {
            match self
                .connection
                .poll()
                .map_err(|e| error!("Polling error for messages: {}", e))?
            {
                Async::Ready(None) => return Ok(false),
                Async::Ready(Some(message)) => {
                    let message_result = self.handle_message(message)?;
                    if !message_result {
                        return Ok(false);
                    }
                }
                Async::NotReady => return Ok(true),
            }
        }
    }
}

impl Future for PubsubConnectionInner {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let flush_req = self.handle_new_subs()?;
        if flush_req {
            self.do_flush()?;
        }
        let cont = self.handle_messages()?;
        if cont {
            Ok(Async::NotReady)
        } else {
            Ok(Async::Ready(()))
        }
    }
}

/// A shareable reference to subscribe to PUBSUB topics
#[derive(Debug, Clone)]
pub struct PubsubConnection {
    out_tx_c:
        Reconnect<PubsubEvent, mpsc::UnboundedSender<PubsubEvent>, error::Error, error::Error>,
}

/// Used for Redis's PUBSUB functionality.
///
/// Returns a future that resolves to a `PubsubConnection`.
pub fn pubsub_connect(
    addr: &SocketAddr,
) -> impl Future<Item = PubsubConnection, Error = error::Error> + Send {
    let addr = *addr;
    future::lazy(move || {
        reconnect(
            |con: &mpsc::UnboundedSender<PubsubEvent>, act| {
                Box::new(future::result(con.unbounded_send(act)).map_err(|e| e.into()))
            },
            move || {
                let con_f = connect(&addr).map_err(|e| e.into()).and_then(|connection| {
                    let (out_tx, out_rx) = mpsc::unbounded();
                    let pubsub_connection_inner =
                        Box::new(PubsubConnectionInner::new(connection, out_rx));
                    let mut default_executor = DefaultExecutor::current();
                    match default_executor.spawn(pubsub_connection_inner) {
                        Ok(_) => Ok(out_tx),
                        Err(e) => Err(error::Error::Internal(format!(
                            "Cannot spawn a pubsub connection: {:?}",
                            e
                        ))),
                    }
                });
                Box::new(con_f)
            },
        )
    })
    .map(|out_tx_c| PubsubConnection { out_tx_c })
    .map_err(|()| error::Error::EndOfStream)
}

impl PubsubConnection {
    /// Subscribes to a particular PUBSUB topic.
    ///
    /// Returns a future that resolves to a `Stream` that contains all the messages published on
    /// that particular topic.
    pub fn subscribe(&self, topic: &str) -> impl Future<Item = PubsubStream, Error = error::Error> {
        let (tx, rx) = mpsc::unbounded();
        let (signal_t, signal_r) = oneshot::channel();
        let do_work_f = self
            .out_tx_c
            .do_work(PubsubEvent::Subscribe(topic.to_owned(), tx, signal_t))
            .map_err(|_| error::Error::EndOfStream);

        let stream = PubsubStream {
            topic: topic.to_owned(),
            underlying: rx,
            con: self.clone(),
        };
        do_work_f.and_then(|()| signal_r.map(|_| stream).map_err(|e| e.into()))
    }

    pub fn unsubscribe<T: Into<String>>(&self, topic: T) {
        // Ignoring any results, as any errors communicating with Redis would de-facto unsubscribe
        // anyway, and would be reported/logged elsewhere
        let _ = self
            .out_tx_c
            .do_work(PubsubEvent::Unsubscribe(topic.into()));
    }
}

pub struct PubsubStream {
    topic: String,
    underlying: PubsubStreamInner,
    con: PubsubConnection,
}

impl Stream for PubsubStream {
    type Item = resp::RespValue;
    type Error = ();

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        self.underlying.poll()
    }
}

impl Drop for PubsubStream {
    fn drop(&mut self) {
        self.con.unsubscribe(self.topic.as_ref());
    }
}
<commit_msg>fix breakage caused by impl From<&String> for String<commit_after>/*
 * Copyright 2017-2019 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
 * http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
 * <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
 * option. This file may not be copied, modified, or distributed
 * except according to those terms.
 */

use std::collections::{hash_map::Entry, HashMap};
use std::net::SocketAddr;

use futures::{
    future,
    stream::Fuse,
    sync::{mpsc, oneshot},
    Async, AsyncSink, Future, Poll, Sink, Stream,
};

use tokio_executor::{DefaultExecutor, Executor};

use super::connect::{connect, RespConnection};
use error;
use reconnect::{reconnect, Reconnect};
use resp::{self, FromResp};

#[derive(Debug)]
enum PubsubEvent {
    Subscribe(String, PubsubSink, oneshot::Sender<()>),
    Unsubscribe(String),
}

pub type PubsubStreamInner = mpsc::UnboundedReceiver<resp::RespValue>;
pub type PubsubSink = mpsc::UnboundedSender<resp::RespValue>;

struct PubsubConnectionInner {
    connection: RespConnection,
    out_rx: Fuse<mpsc::UnboundedReceiver<PubsubEvent>>,
    subscriptions: HashMap<String, PubsubSink>,
    pending_subs: HashMap<String, (PubsubSink, oneshot::Sender<()>)>,
    send_pending: Option<resp::RespValue>,
}

impl PubsubConnectionInner {
    fn new(con: RespConnection, out_rx: mpsc::UnboundedReceiver<PubsubEvent>) -> Self {
        PubsubConnectionInner {
            connection: con,
            out_rx: out_rx.fuse(),
            subscriptions: HashMap::new(),
            pending_subs: HashMap::new(),
            send_pending: None,
        }
    }

    /// Returns true = OK, more can be sent, or false = sink is full, needs flushing
    fn do_send(&mut self, msg: resp::RespValue) -> Result<bool, ()> {
        match self
            .connection
            .start_send(msg)
            .map_err(|e| error!("Cannot send subscription request to Redis: {}", e))?
        {
            AsyncSink::Ready => Ok(true),
            AsyncSink::NotReady(msg) => {
                self.send_pending = Some(msg);
                Ok(false)
            }
        }
    }

    fn do_flush(&mut self) -> Result<(), ()> {
        self.connection
            .poll_complete()
            .map(|_| ())
            .map_err(|e| error!("Error polling for completeness: {}", e))
    }

    // Returns true = flushing required.  false = no flushing required
    fn handle_new_subs(&mut self) -> Result<bool, ()> {
        let mut flushing_req = false;
        if let Some(msg) = self.send_pending.take() {
            flushing_req = true;
            if !self.do_send(msg)? {
                return Ok(flushing_req);
            }
        }
        loop {
            match self
                .out_rx
                .poll()
                .map_err(|_| error!("Cannot poll for new subscriptions"))?
            {
                Async::Ready(None) => {
                    return Ok(flushing_req);
                }
                Async::Ready(Some(pubsub_event)) => {
                    let message = match pubsub_event {
                        PubsubEvent::Subscribe(topic, sender, signal) => {
                            self.pending_subs.insert(topic.clone(), (sender, signal));
                            resp_array!["SUBSCRIBE", topic]
                        }
                        PubsubEvent::Unsubscribe(topic) => resp_array!["UNSUBSCRIBE", topic],
                    };
                    flushing_req = true;
                    if !self.do_send(message)? {
                        return Ok(flushing_req);
                    }
                }
                Async::NotReady => {
                    return Ok(flushing_req);
                }
            }
        }
    }

    fn handle_message(&mut self, msg: resp::RespValue) -> Result<bool, ()> {
        let (message_type, topic, msg) = match msg {
            resp::RespValue::Array(mut messages) => match (
                messages.pop(),
                messages.pop(),
                messages.pop(),
                messages.pop(),
            ) {
                (Some(msg), Some(topic), Some(message_type), None) => {
                    match (msg, String::from_resp(topic), message_type) {
                        (msg, Ok(topic), resp::RespValue::BulkString(bytes)) => (bytes, topic, msg),
                        _ => {
                            error!("Incorrect format of PUBSUB message");
                            return Err(());
                        }
                    }
                }
                _ => {
                    error!("Wrong number of parts for a PUBSUB message");
                    return Err(());
                }
            },
            _ => {
                error!("PUBSUB message should be encoded as an array");
                return Err(());
            }
        };

        if message_type == b"subscribe" {
            if let Some((sender, signal)) = self.pending_subs.remove(&topic) {
                self.subscriptions.insert(topic, sender);
                signal
                    .send(())
                    .map_err(|_| error!("Error confirming subscription"))?;
            }
        } else if message_type == b"unsubscribe" {
            if let Entry::Occupied(entry) = self.subscriptions.entry(topic) {
                entry.remove_entry();
            }
            if self.subscriptions.is_empty() {
                return Ok(false);
            }
        } else if message_type == b"message" {
            if let Some(sender) = self.subscriptions.get(&topic) {
                sender.unbounded_send(msg).expect("Cannot send message");
            }
        }

        Ok(true)
    }

    /// Returns true, if there are still valid subscriptions at the end, or false if not, i.e. the whole thing can be dropped.
    fn handle_messages(&mut self) -> Result<bool, ()> {
        loop {
            match self
                .connection
                .poll()
                .map_err(|e| error!("Polling error for messages: {}", e))?
            {
                Async::Ready(None) => return Ok(false),
                Async::Ready(Some(message)) => {
                    let message_result = self.handle_message(message)?;
                    if !message_result {
                        return Ok(false);
                    }
                }
                Async::NotReady => return Ok(true),
            }
        }
    }
}

impl Future for PubsubConnectionInner {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let flush_req = self.handle_new_subs()?;
        if flush_req {
            self.do_flush()?;
        }
        let cont = self.handle_messages()?;
        if cont {
            Ok(Async::NotReady)
        } else {
            Ok(Async::Ready(()))
        }
    }
}

/// A shareable reference to subscribe to PUBSUB topics
#[derive(Debug, Clone)]
pub struct PubsubConnection {
    out_tx_c:
        Reconnect<PubsubEvent, mpsc::UnboundedSender<PubsubEvent>, error::Error, error::Error>,
}

/// Used for Redis's PUBSUB functionality.
///
/// Returns a future that resolves to a `PubsubConnection`.
pub fn pubsub_connect(
    addr: &SocketAddr,
) -> impl Future<Item = PubsubConnection, Error = error::Error> + Send {
    let addr = *addr;
    future::lazy(move || {
        reconnect(
            |con: &mpsc::UnboundedSender<PubsubEvent>, act| {
                Box::new(future::result(con.unbounded_send(act)).map_err(|e| e.into()))
            },
            move || {
                let con_f = connect(&addr).map_err(|e| e.into()).and_then(|connection| {
                    let (out_tx, out_rx) = mpsc::unbounded();
                    let pubsub_connection_inner =
                        Box::new(PubsubConnectionInner::new(connection, out_rx));
                    let mut default_executor = DefaultExecutor::current();
                    match default_executor.spawn(pubsub_connection_inner) {
                        Ok(_) => Ok(out_tx),
                        Err(e) => Err(error::Error::Internal(format!(
                            "Cannot spawn a pubsub connection: {:?}",
                            e
                        ))),
                    }
                });
                Box::new(con_f)
            },
        )
    })
    .map(|out_tx_c| PubsubConnection { out_tx_c })
    .map_err(|()| error::Error::EndOfStream)
}

impl PubsubConnection {
    /// Subscribes to a particular PUBSUB topic.
    ///
    /// Returns a future that resolves to a `Stream` that contains all the messages published on
    /// that particular topic.
    pub fn subscribe(&self, topic: &str) -> impl Future<Item = PubsubStream, Error = error::Error> {
        let (tx, rx) = mpsc::unbounded();
        let (signal_t, signal_r) = oneshot::channel();
        let do_work_f = self
            .out_tx_c
            .do_work(PubsubEvent::Subscribe(topic.to_owned(), tx, signal_t))
            .map_err(|_| error::Error::EndOfStream);

        let stream = PubsubStream {
            topic: topic.to_owned(),
            underlying: rx,
            con: self.clone(),
        };
        do_work_f.and_then(|()| signal_r.map(|_| stream).map_err(|e| e.into()))
    }

    pub fn unsubscribe<T: Into<String>>(&self, topic: T) {
        // Ignoring any results, as any errors communicating with Redis would de-facto unsubscribe
        // anyway, and would be reported/logged elsewhere
        let _ = self
            .out_tx_c
            .do_work(PubsubEvent::Unsubscribe(topic.into()));
    }
}

pub struct PubsubStream {
    topic: String,
    underlying: PubsubStreamInner,
    con: PubsubConnection,
}

impl Stream for PubsubStream {
    type Item = resp::RespValue;
    type Error = ();

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        self.underlying.poll()
    }
}

impl Drop for PubsubStream {
    fn drop(&mut self) {
        let topic: &str = self.topic.as_ref();
        self.con.unsubscribe(topic);
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use LinkerFlavor;
use std::io;
use std::process::Command;
use target::{LinkArgs, TargetOptions};

use self::Arch::*;

#[allow(non_camel_case_types)]
#[derive(Copy, Clone)]
pub enum Arch {
    Armv7,
    Armv7s,
    Arm64,
    I386,
    X86_64
}

impl Arch {
    pub fn to_string(&self) -> &'static str {
        match self {
            &Armv7 => "armv7",
            &Armv7s => "armv7s",
            &Arm64 => "arm64",
            &I386 => "i386",
            &X86_64 => "x86_64"
        }
    }
}

pub fn get_sdk_root(sdk_name: &str) -> Result<String, String> {
    let res = Command::new("xcrun")
                      .arg("--show-sdk-path")
                      .arg("-sdk")
                      .arg(sdk_name)
                      .output()
                      .and_then(|output| {
                          if output.status.success() {
                              Ok(String::from_utf8(output.stdout).unwrap())
                          } else {
                              let error = String::from_utf8(output.stderr);
                              let error = format!("process exit with error: {}",
                                                  error.unwrap());
                              Err(io::Error::new(io::ErrorKind::Other,
                                                 &error[..]))
                          }
                      });

    match res {
        Ok(output) => Ok(output.trim().to_string()),
        Err(e) => Err(format!("failed to get {} SDK path: {}", sdk_name, e))
    }
}

fn build_pre_link_args(arch: Arch) -> Result<LinkArgs, String> {
    let sdk_name = match arch {
        Armv7 | Armv7s | Arm64 => "iphoneos",
        I386 | X86_64 => "iphonesimulator"
    };

    let arch_name = arch.to_string();

    let sdk_root = get_sdk_root(sdk_name)?;

    let mut args = LinkArgs::new();
    args.insert(LinkerFlavor::Gcc,
                vec!["-arch".to_string(),
                     arch_name.to_string(),
                     "-Wl,-syslibroot".to_string(),
                     sdk_root]);

    Ok(args)
}

fn target_cpu(arch: Arch) -> String {
    match arch {
        Armv7 => "cortex-a8", // iOS7 is supported on iPhone 4 and higher
        Armv7s => "cortex-a9",
        Arm64 => "cyclone",
        I386 => "yonah",
        X86_64 => "core2",
    }.to_string()
}

pub fn opts(arch: Arch) -> Result<TargetOptions, String> {
    let pre_link_args = build_pre_link_args(arch)?;
    Ok(TargetOptions {
        cpu: target_cpu(arch),
        dynamic_linking: false,
        executables: true,
        pre_link_args,
        has_elf_tls: false,
        exe_allocation_crate: None,
        .. super::apple_base::opts()
    })
}
<commit_msg>explanatory note for the workaround<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use LinkerFlavor;
use std::io;
use std::process::Command;
use target::{LinkArgs, TargetOptions};

use self::Arch::*;

#[allow(non_camel_case_types)]
#[derive(Copy, Clone)]
pub enum Arch {
    Armv7,
    Armv7s,
    Arm64,
    I386,
    X86_64
}

impl Arch {
    pub fn to_string(&self) -> &'static str {
        match self {
            &Armv7 => "armv7",
            &Armv7s => "armv7s",
            &Arm64 => "arm64",
            &I386 => "i386",
            &X86_64 => "x86_64"
        }
    }
}

pub fn get_sdk_root(sdk_name: &str) -> Result<String, String> {
    let res = Command::new("xcrun")
                      .arg("--show-sdk-path")
                      .arg("-sdk")
                      .arg(sdk_name)
                      .output()
                      .and_then(|output| {
                          if output.status.success() {
                              Ok(String::from_utf8(output.stdout).unwrap())
                          } else {
                              let error = String::from_utf8(output.stderr);
                              let error = format!("process exit with error: {}",
                                                  error.unwrap());
                              Err(io::Error::new(io::ErrorKind::Other,
                                                 &error[..]))
                          }
                      });

    match res {
        Ok(output) => Ok(output.trim().to_string()),
        Err(e) => Err(format!("failed to get {} SDK path: {}", sdk_name, e))
    }
}

fn build_pre_link_args(arch: Arch) -> Result<LinkArgs, String> {
    let sdk_name = match arch {
        Armv7 | Armv7s | Arm64 => "iphoneos",
        I386 | X86_64 => "iphonesimulator"
    };

    let arch_name = arch.to_string();

    let sdk_root = get_sdk_root(sdk_name)?;

    let mut args = LinkArgs::new();
    args.insert(LinkerFlavor::Gcc,
                vec!["-arch".to_string(),
                     arch_name.to_string(),
                     "-Wl,-syslibroot".to_string(),
                     sdk_root]);

    Ok(args)
}

fn target_cpu(arch: Arch) -> String {
    match arch {
        Armv7 => "cortex-a8", // iOS7 is supported on iPhone 4 and higher
        Armv7s => "cortex-a9",
        Arm64 => "cyclone",
        I386 => "yonah",
        X86_64 => "core2",
    }.to_string()
}

pub fn opts(arch: Arch) -> Result<TargetOptions, String> {
    let pre_link_args = build_pre_link_args(arch)?;
    Ok(TargetOptions {
        cpu: target_cpu(arch),
        dynamic_linking: false,
        executables: true,
        pre_link_args,
        has_elf_tls: false,
        // The following line is a workaround for jemalloc 4.5 being broken on
        // ios. jemalloc 5.0 is supposed to fix this.
        // see https://github.com/rust-lang/rust/issues/45262
        exe_allocation_crate: None,
        .. super::apple_base::opts()
    })
}
<|endoftext|>