<commit_before>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::default::Default;
use std::error;
use std::fmt::{self, Display};
use std::os::unix::io::RawFd;
use std::str::FromStr;

use audio_streams::{
    shm_streams::{NullShmStreamSource, ShmStreamSource},
    StreamEffect,
};
use libcras::CrasClient;
use resources::{Alloc, MmioType, SystemAllocator};
use sys_util::{error, EventFd, GuestMemory};

use crate::pci::ac97_bus_master::Ac97BusMaster;
use crate::pci::ac97_mixer::Ac97Mixer;
use crate::pci::ac97_regs::*;
use crate::pci::pci_configuration::{
    PciBarConfiguration, PciClassCode, PciConfiguration, PciHeaderType, PciMultimediaSubclass,
};
use crate::pci::pci_device::{self, PciDevice, Result};
use crate::pci::PciInterruptPin;

// Use 82801AA because it's what qemu does.
const PCI_DEVICE_ID_INTEL_82801AA_5: u16 = 0x2415;

/// AC97 audio device emulation.
/// Provides the PCI interface for the internal Ac97 emulation.
/// Internally the `Ac97BusMaster` and `Ac97Mixer` structs are used to emulated the bus master and
/// mixer registers respectively. `Ac97BusMaster` handles moving smaples between guest memory and
/// the audio backend.
#[derive(Debug, Clone)]
pub enum Ac97Backend {
    NULL,
    CRAS,
}

impl Default for Ac97Backend {
    fn default() -> Self {
        Ac97Backend::NULL
    }
}

/// Errors that are possible from a `Ac97`.
#[derive(Debug)]
pub enum Ac97Error {
    InvalidBackend,
}

impl error::Error for Ac97Error {}

impl Display for Ac97Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Ac97Error::InvalidBackend => write!(f, "Must be cras or null"),
        }
    }
}

impl FromStr for Ac97Backend {
    type Err = Ac97Error;
    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        match s {
            "cras" => Ok(Ac97Backend::CRAS),
            "null" => Ok(Ac97Backend::NULL),
            _ => Err(Ac97Error::InvalidBackend),
        }
    }
}

/// Holds the parameters for a AC97 device
#[derive(Default, Debug, Clone)]
pub struct Ac97Parameters {
    pub backend: Ac97Backend,
    pub capture: bool,
    pub capture_effects: Vec<StreamEffect>,
}

pub struct Ac97Dev {
    config_regs: PciConfiguration,
    pci_bus_dev: Option<(u8, u8)>,
    // The irq events are temporarily saved here. They need to be passed to the device after the
    // jail forks. This happens when the bus is first written.
    irq_evt: Option<EventFd>,
    irq_resample_evt: Option<EventFd>,
    bus_master: Ac97BusMaster,
    mixer: Ac97Mixer,
}

impl Ac97Dev {
    /// Creates an 'Ac97Dev' that uses the given `GuestMemory` and starts with all registers at
    /// default values.
    pub fn new(mem: GuestMemory, audio_server: Box<dyn ShmStreamSource>) -> Self {
        let config_regs = PciConfiguration::new(
            0x8086,
            PCI_DEVICE_ID_INTEL_82801AA_5,
            PciClassCode::MultimediaController,
            &PciMultimediaSubclass::AudioDevice,
            None, // No Programming interface.
            PciHeaderType::Device,
            0x8086, // Subsystem Vendor ID
            0x1,    // Subsystem ID.
        );

        Ac97Dev {
            config_regs,
            pci_bus_dev: None,
            irq_evt: None,
            irq_resample_evt: None,
            bus_master: Ac97BusMaster::new(mem, audio_server),
            mixer: Ac97Mixer::new(),
        }
    }

    fn create_cras_audio_device(params: Ac97Parameters, mem: GuestMemory) -> Result<Ac97Dev> {
        let mut server =
            Box::new(CrasClient::new().map_err(|e| pci_device::Error::CreateCrasClientFailed(e))?);
        if params.capture {
            server.enable_cras_capture();
        }

        let mut cras_audio = Ac97Dev::new(mem, server);
        cras_audio.set_capture_effects(params.capture_effects);
        Ok(cras_audio)
    }

    fn create_null_audio_device(mem: GuestMemory) -> Result<Ac97Dev> {
        let server = Box::new(NullShmStreamSource::new());
        let null_audio = Ac97Dev::new(mem, server);
        Ok(null_audio)
    }

    /// Creates an 'Ac97Dev' with suitable audio server inside based on Ac97Parameters
    pub fn try_new(mem: GuestMemory, param: Ac97Parameters) -> Result<Ac97Dev> {
        match param.backend {
            Ac97Backend::CRAS => Ac97Dev::create_cras_audio_device(param, mem),
            Ac97Backend::NULL => Ac97Dev::create_null_audio_device(mem),
        }
    }

    /// Provides the effect needed in capture stream creation
    pub fn set_capture_effects(&mut self, effect: Vec<StreamEffect>) {
        self.bus_master.set_capture_effects(effect);
    }

    fn read_mixer(&mut self, offset: u64, data: &mut [u8]) {
        match data.len() {
            // The mixer is only accessed with 16-bit words.
            2 => {
                let val: u16 = self.mixer.readw(offset);
                data[0] = val as u8;
                data[1] = (val >> 8) as u8;
            }
            l => error!("mixer read length of {}", l),
        }
    }

    fn write_mixer(&mut self, offset: u64, data: &[u8]) {
        match data.len() {
            // The mixer is only accessed with 16-bit words.
            2 => self
                .mixer
                .writew(offset, u16::from(data[0]) | u16::from(data[1]) << 8),
            l => error!("mixer write length of {}", l),
        }
        // Apply the new mixer settings to the bus master.
        self.bus_master.update_mixer_settings(&self.mixer);
    }

    fn read_bus_master(&mut self, offset: u64, data: &mut [u8]) {
        match data.len() {
            1 => data[0] = self.bus_master.readb(offset),
            2 => {
                let val: u16 = self.bus_master.readw(offset);
                data[0] = val as u8;
                data[1] = (val >> 8) as u8;
            }
            4 => {
                let val: u32 = self.bus_master.readl(offset);
                data[0] = val as u8;
                data[1] = (val >> 8) as u8;
                data[2] = (val >> 16) as u8;
                data[3] = (val >> 24) as u8;
            }
            l => error!("read length of {}", l),
        }
    }

    fn write_bus_master(&mut self, offset: u64, data: &[u8]) {
        match data.len() {
            1 => self.bus_master.writeb(offset, data[0], &self.mixer),
            2 => self
                .bus_master
                .writew(offset, u16::from(data[0]) | u16::from(data[1]) << 8),
            4 => self.bus_master.writel(
                offset,
                (u32::from(data[0]))
                    | (u32::from(data[1]) << 8)
                    | (u32::from(data[2]) << 16)
                    | (u32::from(data[3]) << 24),
            ),
            l => error!("write length of {}", l),
        }
    }
}

impl PciDevice for Ac97Dev {
    fn debug_label(&self) -> String {
        "AC97".to_owned()
    }

    fn assign_bus_dev(&mut self, bus: u8, device: u8) {
        self.pci_bus_dev = Some((bus, device));
    }

    fn assign_irq(
        &mut self,
        irq_evt: EventFd,
        irq_resample_evt: EventFd,
        irq_num: u32,
        irq_pin: PciInterruptPin,
    ) {
        self.config_regs.set_irq(irq_num as u8, irq_pin);
        self.irq_evt = Some(irq_evt);
        self.irq_resample_evt = Some(irq_resample_evt);
    }

    fn allocate_io_bars(&mut self, resources: &mut SystemAllocator) -> Result<Vec<(u64, u64)>> {
        let (bus, dev) = self
            .pci_bus_dev
            .expect("assign_bus_dev must be called prior to allocate_io_bars");
        let mut ranges = Vec::new();
        let mixer_regs_addr = resources
            .mmio_allocator(MmioType::Low)
            .allocate_with_align(
                MIXER_REGS_SIZE,
                Alloc::PciBar { bus, dev, bar: 0 },
                "ac97-mixer_regs".to_string(),
                MIXER_REGS_SIZE,
            )
            .map_err(|e| pci_device::Error::IoAllocationFailed(MIXER_REGS_SIZE, e))?;
        let mixer_config = PciBarConfiguration::default()
            .set_register_index(0)
            .set_address(mixer_regs_addr)
            .set_size(MIXER_REGS_SIZE);
        self.config_regs
            .add_pci_bar(&mixer_config)
            .map_err(|e| pci_device::Error::IoRegistrationFailed(mixer_regs_addr, e))?;
        ranges.push((mixer_regs_addr, MIXER_REGS_SIZE));

        let master_regs_addr = resources
            .mmio_allocator(MmioType::Low)
            .allocate_with_align(
                MASTER_REGS_SIZE,
                Alloc::PciBar { bus, dev, bar: 1 },
                "ac97-master_regs".to_string(),
                MASTER_REGS_SIZE,
            )
            .map_err(|e| pci_device::Error::IoAllocationFailed(MASTER_REGS_SIZE, e))?;
        let master_config = PciBarConfiguration::default()
            .set_register_index(1)
            .set_address(master_regs_addr)
            .set_size(MASTER_REGS_SIZE);
        self.config_regs
            .add_pci_bar(&master_config)
            .map_err(|e| pci_device::Error::IoRegistrationFailed(master_regs_addr, e))?;
        ranges.push((master_regs_addr, MASTER_REGS_SIZE));
        Ok(ranges)
    }

    fn read_config_register(&self, reg_idx: usize) -> u32 {
        self.config_regs.read_reg(reg_idx)
    }

    fn write_config_register(&mut self, reg_idx: usize, offset: u64, data: &[u8]) {
        (&mut self.config_regs).write_reg(reg_idx, offset, data)
    }

    fn keep_fds(&self) -> Vec<RawFd> {
        if let Some(server_fds) = self.bus_master.keep_fds() {
            server_fds
        } else {
            Vec::new()
        }
    }

    fn read_bar(&mut self, addr: u64, data: &mut [u8]) {
        let bar0 = self.config_regs.get_bar_addr(0);
        let bar1 = self.config_regs.get_bar_addr(1);
        match addr {
            a if a >= bar0 && a < bar0 + MIXER_REGS_SIZE => self.read_mixer(addr - bar0, data),
            a if a >= bar1 && a < bar1 + MASTER_REGS_SIZE => {
                self.read_bus_master(addr - bar1, data)
            }
            _ => (),
        }
    }

    fn write_bar(&mut self, addr: u64, data: &[u8]) {
        let bar0 = self.config_regs.get_bar_addr(0);
        let bar1 = self.config_regs.get_bar_addr(1);
        match addr {
            a if a >= bar0 && a < bar0 + MIXER_REGS_SIZE => self.write_mixer(addr - bar0, data),
            a if a >= bar1 && a < bar1 + MASTER_REGS_SIZE => {
                // Check if the irq needs to be passed to the device.
                if let (Some(irq_evt), Some(irq_resample_evt)) =
                    (self.irq_evt.take(), self.irq_resample_evt.take())
                {
                    self.bus_master.set_irq_event_fd(irq_evt, irq_resample_evt);
                }
                self.write_bus_master(addr - bar1, data)
            }
            _ => (),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use audio_streams::shm_streams::MockShmStreamSource;
    use sys_util::GuestAddress;

    #[test]
    fn create() {
        let mem = GuestMemory::new(&[(GuestAddress(0u64), 4 * 1024 * 1024)]).unwrap();
        let mut ac97_dev = Ac97Dev::new(mem, Box::new(MockShmStreamSource::new()));
        let mut allocator = SystemAllocator::builder()
            .add_io_addresses(0x1000_0000, 0x1000_0000)
            .add_low_mmio_addresses(0x2000_0000, 0x1000_0000)
            .add_high_mmio_addresses(0x3000_0000, 0x1000_0000)
            .create_allocator(5, false)
            .unwrap();
        ac97_dev.assign_bus_dev(0, 0);
        assert!(ac97_dev.allocate_io_bars(&mut allocator).is_ok());
    }
}
<commit_msg>ac97: Set CRAS_CLIENT_TYPE_CROSVM for crosvm<commit_after>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::default::Default;
use std::error;
use std::fmt::{self, Display};
use std::os::unix::io::RawFd;
use std::str::FromStr;

use audio_streams::{
    shm_streams::{NullShmStreamSource, ShmStreamSource},
    StreamEffect,
};
use libcras::{CrasClient, CrasClientType};
use resources::{Alloc, MmioType, SystemAllocator};
use sys_util::{error, EventFd, GuestMemory};

use crate::pci::ac97_bus_master::Ac97BusMaster;
use crate::pci::ac97_mixer::Ac97Mixer;
use crate::pci::ac97_regs::*;
use crate::pci::pci_configuration::{
    PciBarConfiguration, PciClassCode, PciConfiguration, PciHeaderType, PciMultimediaSubclass,
};
use crate::pci::pci_device::{self, PciDevice, Result};
use crate::pci::PciInterruptPin;

// Use 82801AA because it's what qemu does.
const PCI_DEVICE_ID_INTEL_82801AA_5: u16 = 0x2415;

/// AC97 audio device emulation.
/// Provides the PCI interface for the internal Ac97 emulation.
/// Internally the `Ac97BusMaster` and `Ac97Mixer` structs are used to emulated the bus master and
/// mixer registers respectively. `Ac97BusMaster` handles moving smaples between guest memory and
/// the audio backend.
#[derive(Debug, Clone)]
pub enum Ac97Backend {
    NULL,
    CRAS,
}

impl Default for Ac97Backend {
    fn default() -> Self {
        Ac97Backend::NULL
    }
}

/// Errors that are possible from a `Ac97`.
#[derive(Debug)]
pub enum Ac97Error {
    InvalidBackend,
}

impl error::Error for Ac97Error {}

impl Display for Ac97Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Ac97Error::InvalidBackend => write!(f, "Must be cras or null"),
        }
    }
}

impl FromStr for Ac97Backend {
    type Err = Ac97Error;
    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        match s {
            "cras" => Ok(Ac97Backend::CRAS),
            "null" => Ok(Ac97Backend::NULL),
            _ => Err(Ac97Error::InvalidBackend),
        }
    }
}

/// Holds the parameters for a AC97 device
#[derive(Default, Debug, Clone)]
pub struct Ac97Parameters {
    pub backend: Ac97Backend,
    pub capture: bool,
    pub capture_effects: Vec<StreamEffect>,
}

pub struct Ac97Dev {
    config_regs: PciConfiguration,
    pci_bus_dev: Option<(u8, u8)>,
    // The irq events are temporarily saved here. They need to be passed to the device after the
    // jail forks. This happens when the bus is first written.
    irq_evt: Option<EventFd>,
    irq_resample_evt: Option<EventFd>,
    bus_master: Ac97BusMaster,
    mixer: Ac97Mixer,
}

impl Ac97Dev {
    /// Creates an 'Ac97Dev' that uses the given `GuestMemory` and starts with all registers at
    /// default values.
    pub fn new(mem: GuestMemory, audio_server: Box<dyn ShmStreamSource>) -> Self {
        let config_regs = PciConfiguration::new(
            0x8086,
            PCI_DEVICE_ID_INTEL_82801AA_5,
            PciClassCode::MultimediaController,
            &PciMultimediaSubclass::AudioDevice,
            None, // No Programming interface.
            PciHeaderType::Device,
            0x8086, // Subsystem Vendor ID
            0x1,    // Subsystem ID.
        );

        Ac97Dev {
            config_regs,
            pci_bus_dev: None,
            irq_evt: None,
            irq_resample_evt: None,
            bus_master: Ac97BusMaster::new(mem, audio_server),
            mixer: Ac97Mixer::new(),
        }
    }

    fn create_cras_audio_device(params: Ac97Parameters, mem: GuestMemory) -> Result<Ac97Dev> {
        let mut server =
            Box::new(CrasClient::new().map_err(|e| pci_device::Error::CreateCrasClientFailed(e))?);
        server.set_client_type(CrasClientType::CRAS_CLIENT_TYPE_CROSVM);
        if params.capture {
            server.enable_cras_capture();
        }

        let mut cras_audio = Ac97Dev::new(mem, server);
        cras_audio.set_capture_effects(params.capture_effects);
        Ok(cras_audio)
    }

    fn create_null_audio_device(mem: GuestMemory) -> Result<Ac97Dev> {
        let server = Box::new(NullShmStreamSource::new());
        let null_audio = Ac97Dev::new(mem, server);
        Ok(null_audio)
    }

    /// Creates an 'Ac97Dev' with suitable audio server inside based on Ac97Parameters
    pub fn try_new(mem: GuestMemory, param: Ac97Parameters) -> Result<Ac97Dev> {
        match param.backend {
            Ac97Backend::CRAS => Ac97Dev::create_cras_audio_device(param, mem),
            Ac97Backend::NULL => Ac97Dev::create_null_audio_device(mem),
        }
    }

    /// Provides the effect needed in capture stream creation
    pub fn set_capture_effects(&mut self, effect: Vec<StreamEffect>) {
        self.bus_master.set_capture_effects(effect);
    }

    fn read_mixer(&mut self, offset: u64, data: &mut [u8]) {
        match data.len() {
            // The mixer is only accessed with 16-bit words.
            2 => {
                let val: u16 = self.mixer.readw(offset);
                data[0] = val as u8;
                data[1] = (val >> 8) as u8;
            }
            l => error!("mixer read length of {}", l),
        }
    }

    fn write_mixer(&mut self, offset: u64, data: &[u8]) {
        match data.len() {
            // The mixer is only accessed with 16-bit words.
            2 => self
                .mixer
                .writew(offset, u16::from(data[0]) | u16::from(data[1]) << 8),
            l => error!("mixer write length of {}", l),
        }
        // Apply the new mixer settings to the bus master.
        self.bus_master.update_mixer_settings(&self.mixer);
    }

    fn read_bus_master(&mut self, offset: u64, data: &mut [u8]) {
        match data.len() {
            1 => data[0] = self.bus_master.readb(offset),
            2 => {
                let val: u16 = self.bus_master.readw(offset);
                data[0] = val as u8;
                data[1] = (val >> 8) as u8;
            }
            4 => {
                let val: u32 = self.bus_master.readl(offset);
                data[0] = val as u8;
                data[1] = (val >> 8) as u8;
                data[2] = (val >> 16) as u8;
                data[3] = (val >> 24) as u8;
            }
            l => error!("read length of {}", l),
        }
    }

    fn write_bus_master(&mut self, offset: u64, data: &[u8]) {
        match data.len() {
            1 => self.bus_master.writeb(offset, data[0], &self.mixer),
            2 => self
                .bus_master
                .writew(offset, u16::from(data[0]) | u16::from(data[1]) << 8),
            4 => self.bus_master.writel(
                offset,
                (u32::from(data[0]))
                    | (u32::from(data[1]) << 8)
                    | (u32::from(data[2]) << 16)
                    | (u32::from(data[3]) << 24),
            ),
            l => error!("write length of {}", l),
        }
    }
}

impl PciDevice for Ac97Dev {
    fn debug_label(&self) -> String {
        "AC97".to_owned()
    }

    fn assign_bus_dev(&mut self, bus: u8, device: u8) {
        self.pci_bus_dev = Some((bus, device));
    }

    fn assign_irq(
        &mut self,
        irq_evt: EventFd,
        irq_resample_evt: EventFd,
        irq_num: u32,
        irq_pin: PciInterruptPin,
    ) {
        self.config_regs.set_irq(irq_num as u8, irq_pin);
        self.irq_evt = Some(irq_evt);
        self.irq_resample_evt = Some(irq_resample_evt);
    }

    fn allocate_io_bars(&mut self, resources: &mut SystemAllocator) -> Result<Vec<(u64, u64)>> {
        let (bus, dev) = self
            .pci_bus_dev
            .expect("assign_bus_dev must be called prior to allocate_io_bars");
        let mut ranges = Vec::new();
        let mixer_regs_addr = resources
            .mmio_allocator(MmioType::Low)
            .allocate_with_align(
                MIXER_REGS_SIZE,
                Alloc::PciBar { bus, dev, bar: 0 },
                "ac97-mixer_regs".to_string(),
                MIXER_REGS_SIZE,
            )
            .map_err(|e| pci_device::Error::IoAllocationFailed(MIXER_REGS_SIZE, e))?;
        let mixer_config = PciBarConfiguration::default()
            .set_register_index(0)
            .set_address(mixer_regs_addr)
            .set_size(MIXER_REGS_SIZE);
        self.config_regs
            .add_pci_bar(&mixer_config)
            .map_err(|e| pci_device::Error::IoRegistrationFailed(mixer_regs_addr, e))?;
        ranges.push((mixer_regs_addr, MIXER_REGS_SIZE));

        let master_regs_addr = resources
            .mmio_allocator(MmioType::Low)
            .allocate_with_align(
                MASTER_REGS_SIZE,
                Alloc::PciBar { bus, dev, bar: 1 },
                "ac97-master_regs".to_string(),
                MASTER_REGS_SIZE,
            )
            .map_err(|e| pci_device::Error::IoAllocationFailed(MASTER_REGS_SIZE, e))?;
        let master_config = PciBarConfiguration::default()
            .set_register_index(1)
            .set_address(master_regs_addr)
            .set_size(MASTER_REGS_SIZE);
        self.config_regs
            .add_pci_bar(&master_config)
            .map_err(|e| pci_device::Error::IoRegistrationFailed(master_regs_addr, e))?;
        ranges.push((master_regs_addr, MASTER_REGS_SIZE));
        Ok(ranges)
    }

    fn read_config_register(&self, reg_idx: usize) -> u32 {
        self.config_regs.read_reg(reg_idx)
    }

    fn write_config_register(&mut self, reg_idx: usize, offset: u64, data: &[u8]) {
        (&mut self.config_regs).write_reg(reg_idx, offset, data)
    }

    fn keep_fds(&self) -> Vec<RawFd> {
        if let Some(server_fds) = self.bus_master.keep_fds() {
            server_fds
        } else {
            Vec::new()
        }
    }

    fn read_bar(&mut self, addr: u64, data: &mut [u8]) {
        let bar0 = self.config_regs.get_bar_addr(0);
        let bar1 = self.config_regs.get_bar_addr(1);
        match addr {
            a if a >= bar0 && a < bar0 + MIXER_REGS_SIZE => self.read_mixer(addr - bar0, data),
            a if a >= bar1 && a < bar1 + MASTER_REGS_SIZE => {
                self.read_bus_master(addr - bar1, data)
            }
            _ => (),
        }
    }

    fn write_bar(&mut self, addr: u64, data: &[u8]) {
        let bar0 = self.config_regs.get_bar_addr(0);
        let bar1 = self.config_regs.get_bar_addr(1);
        match addr {
            a if a >= bar0 && a < bar0 + MIXER_REGS_SIZE => self.write_mixer(addr - bar0, data),
            a if a >= bar1 && a < bar1 + MASTER_REGS_SIZE => {
                // Check if the irq needs to be passed to the device.
                if let (Some(irq_evt), Some(irq_resample_evt)) =
                    (self.irq_evt.take(), self.irq_resample_evt.take())
                {
                    self.bus_master.set_irq_event_fd(irq_evt, irq_resample_evt);
                }
                self.write_bus_master(addr - bar1, data)
            }
            _ => (),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use audio_streams::shm_streams::MockShmStreamSource;
    use sys_util::GuestAddress;

    #[test]
    fn create() {
        let mem = GuestMemory::new(&[(GuestAddress(0u64), 4 * 1024 * 1024)]).unwrap();
        let mut ac97_dev = Ac97Dev::new(mem, Box::new(MockShmStreamSource::new()));
        let mut allocator = SystemAllocator::builder()
            .add_io_addresses(0x1000_0000, 0x1000_0000)
            .add_low_mmio_addresses(0x2000_0000, 0x1000_0000)
            .add_high_mmio_addresses(0x3000_0000, 0x1000_0000)
            .create_allocator(5, false)
            .unwrap();
        ac97_dev.assign_bus_dev(0, 0);
        assert!(ac97_dev.allocate_io_bars(&mut allocator).is_ok());
    }
}
<|endoftext|>


<commit_before>#[cfg(target_os = "freebsd")]
use std::ffi::c_void;
#[cfg(target_os = "macos")]
use std::ffi::CString;
#[cfg(target_os = "macos")]
use std::iter::repeat;
#[cfg(target_os = "freebsd")]
use std::iter::repeat;
#[cfg(target_os = "freebsd")]
use std::mem::size_of;
#[cfg(target_os = "macos")]
use std::ptr::null_mut;
use std::{
    fs::File,
    io::{Seek, SeekFrom},
    os::unix::io::{AsRawFd, FromRawFd, IntoRawFd},
    os::unix::process::CommandExt,
    process::Command,
};

use anyhow::{bail, Context};
use libc::{self, pid_t};
#[cfg(target_os = "macos")]
use libc::{c_char, PATH_MAX};
#[cfg(target_os = "freebsd")]
use libc::{sysctl, CTL_KERN, KERN_PROC, KERN_PROC_PATHNAME, PATH_MAX};
use mio::net::UnixStream;
use nix::{self, unistd::*};
use serde_json;
use tempfile::tempfile;

use sozu::{metrics, server::Server};
use sozu_command_lib::{
    channel::Channel,
    config::Config,
    logging::target_to_backend,
    proxy::{ProxyRequest, ProxyRequestData, ProxyResponse},
    ready::Ready,
    scm_socket::{Listeners, ScmSocket},
    state::ConfigState,
};

use crate::{command::Worker, logging, util};

pub fn start_workers(executable_path: String, config: &Config) -> anyhow::Result<Vec<Worker>> {
    let state = ConfigState::new();
    let mut workers = Vec::new();
    for index in 0..config.worker_count {
        let listeners = Some(Listeners {
            http: Vec::new(),
            tls: Vec::new(),
            tcp: Vec::new(),
        });

        let (pid, command_channel, scm_socket) = start_worker_process(
            &index.to_string(),
            config,
            executable_path.clone(),
            &state,
            listeners,
        )?;
        let mut worker = Worker::new(index as u32, pid, command_channel, scm_socket, config);

        // the new worker expects a status message at startup
        if let Some(command_channel) = worker.command_channel.as_mut() {
            command_channel.set_blocking(true);
            command_channel.write_message(&ProxyRequest {
                id: format!("start-status-{}", index),
                order: ProxyRequestData::Status,
            });
            command_channel.set_nonblocking(true);
        }

        workers.push(worker);
    }
    info!("Created workers");
    Ok(workers)
}

pub fn start_worker(
    id: u32,
    config: &Config,
    executable_path: String,
    state: &ConfigState,
    listeners: Option<Listeners>,
) -> anyhow::Result<Worker> {
    let (pid, command, scm) =
        start_worker_process(&id.to_string(), config, executable_path, state, listeners)?;

    Ok(Worker::new(id, pid, command, scm, config))
}

pub fn begin_worker_process(
    command_socket_fd: i32,
    scm_socket_fd: i32,
    configuration_state_fd: i32,
    id: i32,
    command_buffer_size: usize,
    max_command_buffer_size: usize,
) -> Result<(), anyhow::Error> {
    let mut command_channel: Channel<ProxyResponse, Config> = Channel::new(
        unsafe { UnixStream::from_raw_fd(command_socket_fd) },
        command_buffer_size,
        max_command_buffer_size,
    );

    command_channel.set_nonblocking(false);

    let configuration_state_file = unsafe { File::from_raw_fd(configuration_state_fd) };
    let config_state: ConfigState = serde_json::from_reader(configuration_state_file)
        .with_context(|| "could not parse configuration state data")?;

    let worker_config = command_channel
        .read_message()
        .with_context(|| "worker could not read configuration from socket")?;
    //println!("got message: {:?}", worker_config);

    let worker_id = format!("{}-{:02}", "WRK", id);
    logging::setup(
        worker_id.clone(),
        &worker_config.log_level,
        &worker_config.log_target,
        worker_config.log_access_target.as_deref(),
    );

    // debug!("Creating worker {} with config: {:#?}", worker_id, worker_config);

    let backend = target_to_backend(&worker_config.log_target);
    let access_backend = worker_config
        .log_access_target
        .as_deref()
        .map(target_to_backend);
    sozu_command_lib::logging::Logger::init(
        worker_id.clone(),
        &worker_config.log_level,
        backend,
        access_backend,
    );
    info!("worker {} starting...", id);

    command_channel.set_nonblocking(true);
    let mut command_channel: Channel<ProxyResponse, ProxyRequest> = command_channel.into();
    command_channel.readiness.insert(Ready::readable());

    if let Some(metrics) = worker_config.metrics.as_ref() {
        metrics::setup(
            &metrics.address,
            worker_id,
            metrics.tagged_metrics,
            metrics.prefix.clone(),
        );
    }

    let mut server = Server::new_from_config(
        command_channel,
        ScmSocket::new(scm_socket_fd),
        worker_config,
        config_state,
        true,
    );

    info!("starting event loop");
    server.run();
    info!("ending event loop");
    Ok(())
}

pub fn start_worker_process(
    id: &str,
    config: &Config,
    executable_path: String,
    state: &ConfigState,
    listeners: Option<Listeners>,
) -> anyhow::Result<(pid_t, Channel<ProxyRequest, ProxyResponse>, ScmSocket)> {
    trace!("parent({})", unsafe { libc::getpid() });

    let mut state_file =
        tempfile().expect("could not create temporary file for configuration state");
    util::disable_close_on_exec(state_file.as_raw_fd())?;

    serde_json::to_writer(&mut state_file, state)
        .with_context(|| "could not write upgrade data to temporary file")?;
    state_file
        .seek(SeekFrom::Start(0))
        .with_context(|| "could not seek to beginning of file")?;

    let (command_socket_main, command_socket_worker) = UnixStream::pair()?;
    let (scm_socket_main, scm_socket_worker) = UnixStream::pair()?;

    let scm_main = ScmSocket::new(scm_socket_main.into_raw_fd());

    util::disable_close_on_exec(command_socket_worker.as_raw_fd())?;
    util::disable_close_on_exec(scm_socket_worker.as_raw_fd())?;

    let mut command_channel: Channel<Config, ProxyResponse> = Channel::new(
        command_socket_main,
        config.command_buffer_size,
        config.max_command_buffer_size,
    );
    command_channel.set_nonblocking(false);

    info!("{} launching worker", id);
    debug!("executable path is {}", executable_path);
    match unsafe { fork() } {
        Ok(ForkResult::Parent { child }) => {
            info!("{} worker launched: {}", id, child);
            command_channel.write_message(config);
            command_channel.set_nonblocking(true);

            if let Some(listeners) = listeners {
                info!("sending listeners to new worker: {:?}", listeners);
                let result = scm_main.send_listeners(&listeners);
                info!("sent listeners from main: {:?}", result);
                listeners.close();
            };
            util::disable_close_on_exec(scm_main.fd)?;

            Ok((child.into(), command_channel.into(), scm_main))
        }
        Ok(ForkResult::Child) => {
            trace!("child({}):\twill spawn a child", unsafe { libc::getpid() });
            Command::new(executable_path)
                .arg("worker")
                .arg("--id")
                .arg(id)
                .arg("--fd")
                .arg(command_socket_worker.as_raw_fd().to_string())
                .arg("--scm")
                .arg(scm_socket_worker.as_raw_fd().to_string())
                .arg("--configuration-state-fd")
                .arg(state_file.as_raw_fd().to_string())
                .arg("--command-buffer-size")
                .arg(config.command_buffer_size.to_string())
                .arg("--max-command-buffer-size")
                .arg(config.max_command_buffer_size.to_string())
                .exec();

            unreachable!();
        }
        Err(e) => {
            error!("Error during fork(): {}", e);
            Err(anyhow::Error::from(e))
        }
    }
}

#[cfg(target_os = "linux")]
pub unsafe fn get_executable_path() -> anyhow::Result<String> {
    use std::fs;

    let path = fs::read_link("/proc/self/exe").with_context(|| "/proc/self/exe doesn't exist")?;

    let mut path_str = match path.into_os_string().into_string() {
        Ok(s) => s,
        Err(_) => bail!("Failed to convert PathBuf to String"),
    };

    if path_str.ends_with(" (deleted)") {
        // The kernel appends " (deleted)" to the symlink when the original executable has been replaced
        let len = path_str.len();
        path_str.truncate(len - 10)
    }

    Ok(path_str)
}

#[cfg(target_os = "macos")]
extern "C" {
    pub fn _NSGetExecutablePath(buf: *mut c_char, size: *mut u32) -> i32;
}

#[cfg(target_os = "macos")]
pub unsafe fn get_executable_path() -> String {
    let capacity = PATH_MAX as usize;
    let mut temp: Vec<u8> = Vec::with_capacity(capacity);
    temp.extend(repeat(0).take(capacity));
    let pathbuf = CString::from_vec_unchecked(temp);
    let ptr = pathbuf.into_raw();

    let mut size = capacity as u32;
    if _NSGetExecutablePath(ptr, &mut size) == 0 {
        let mut temp2: Vec<u8> = Vec::with_capacity(capacity);
        temp2.extend(repeat(0).take(capacity));
        let pathbuf2 = CString::from_vec_unchecked(temp2);
        let ptr2 = pathbuf2.into_raw();

        if libc::realpath(ptr, ptr2) != null_mut() {
            let path = CString::from_raw(ptr2);
            path.to_str()
                .expect("failed to convert CString to String")
                .to_string()
        } else {
            panic!();
        }
    } else {
        panic!("buffer too small");
    }
}

#[cfg(target_os = "freebsd")]
pub unsafe fn get_executable_path() -> String {
    let mut capacity = PATH_MAX as usize;
    let mut path: Vec<u8> = Vec::with_capacity(capacity);
    path.extend(repeat(0).take(capacity));

    let mib: Vec<i32> = vec![CTL_KERN, KERN_PROC, KERN_PROC_PATHNAME, -1];
    let len = mib.len() * size_of::<i32>();
    let element_size = size_of::<i32>();

    let res = sysctl(
        mib.as_ptr(),
        (len / element_size) as u32,
        path.as_mut_ptr() as *mut c_void,
        &mut capacity,
        std::ptr::null() as *const c_void,
        0,
    );
    if res != 0 {
        panic!("Could not retrieve the path of the executable");
    }

    String::from_raw_parts(path.as_mut_ptr(), capacity - 1, path.len())
}
<commit_msg>rename start_worker_process into fork_main_into_worker<commit_after>#[cfg(target_os = "freebsd")]
use std::ffi::c_void;
#[cfg(target_os = "macos")]
use std::ffi::CString;
#[cfg(target_os = "macos")]
use std::iter::repeat;
#[cfg(target_os = "freebsd")]
use std::iter::repeat;
#[cfg(target_os = "freebsd")]
use std::mem::size_of;
#[cfg(target_os = "macos")]
use std::ptr::null_mut;
use std::{
    fs::File,
    io::{Seek, SeekFrom},
    os::unix::io::{AsRawFd, FromRawFd, IntoRawFd},
    os::unix::process::CommandExt,
    process::Command,
};

use anyhow::{bail, Context};
use libc::{self, pid_t};
#[cfg(target_os = "macos")]
use libc::{c_char, PATH_MAX};
#[cfg(target_os = "freebsd")]
use libc::{sysctl, CTL_KERN, KERN_PROC, KERN_PROC_PATHNAME, PATH_MAX};
use mio::net::UnixStream;
use nix::{self, unistd::*};
use serde_json;
use tempfile::tempfile;

use sozu::{metrics, server::Server};
use sozu_command_lib::{
    channel::Channel,
    config::Config,
    logging::target_to_backend,
    proxy::{ProxyRequest, ProxyRequestData, ProxyResponse},
    ready::Ready,
    scm_socket::{Listeners, ScmSocket},
    state::ConfigState,
};

use crate::{command::Worker, logging, util};

pub fn start_workers(executable_path: String, config: &Config) -> anyhow::Result<Vec<Worker>> {
    let state = ConfigState::new();
    let mut workers = Vec::new();
    for index in 0..config.worker_count {
        let listeners = Some(Listeners {
            http: Vec::new(),
            tls: Vec::new(),
            tcp: Vec::new(),
        });

        let (pid, command_channel, scm_socket) = fork_main_into_worker(
            &index.to_string(),
            config,
            executable_path.clone(),
            &state,
            listeners,
        )?;
        let mut worker = Worker::new(index as u32, pid, command_channel, scm_socket, config);

        // the new worker expects a status message at startup
        if let Some(command_channel) = worker.command_channel.as_mut() {
            command_channel.set_blocking(true);
            command_channel.write_message(&ProxyRequest {
                id: format!("start-status-{}", index),
                order: ProxyRequestData::Status,
            });
            command_channel.set_nonblocking(true);
        }

        workers.push(worker);
    }
    info!("Created workers");
    Ok(workers)
}

/// called by the CommandServer to start an individual worker
pub fn start_worker(
    id: u32,
    config: &Config,
    executable_path: String,
    state: &ConfigState,
    listeners: Option<Listeners>,
) -> anyhow::Result<Worker> {
    let (worker_pid, command_channel, scm_socket) =
        fork_main_into_worker(&id.to_string(), config, executable_path, state, listeners)?;

    Ok(Worker::new(
        id,
        worker_pid,
        command_channel,
        scm_socket,
        config,
    ))
}

/// called within a worker process, this starts the actual proxy
pub fn begin_worker_process(
    command_socket_fd: i32,
    scm_socket_fd: i32,
    configuration_state_fd: i32,
    id: i32,
    command_buffer_size: usize,
    max_command_buffer_size: usize,
) -> Result<(), anyhow::Error> {
    let mut command_channel: Channel<ProxyResponse, Config> = Channel::new(
        unsafe { UnixStream::from_raw_fd(command_socket_fd) },
        command_buffer_size,
        max_command_buffer_size,
    );

    command_channel.set_nonblocking(false);

    let configuration_state_file = unsafe { File::from_raw_fd(configuration_state_fd) };
    let config_state: ConfigState = serde_json::from_reader(configuration_state_file)
        .with_context(|| "could not parse configuration state data")?;

    let worker_config = command_channel
        .read_message()
        .with_context(|| "worker could not read configuration from socket")?;
    //println!("got message: {:?}", worker_config);

    let worker_id = format!("{}-{:02}", "WRK", id);
    logging::setup(
        worker_id.clone(),
        &worker_config.log_level,
        &worker_config.log_target,
        worker_config.log_access_target.as_deref(),
    );

    // debug!("Creating worker {} with config: {:#?}", worker_id, worker_config);

    let backend = target_to_backend(&worker_config.log_target);
    let access_backend = worker_config
        .log_access_target
        .as_deref()
        .map(target_to_backend);
    sozu_command_lib::logging::Logger::init(
        worker_id.clone(),
        &worker_config.log_level,
        backend,
        access_backend,
    );
    info!("worker {} starting...", id);

    command_channel.set_nonblocking(true);
    let mut command_channel: Channel<ProxyResponse, ProxyRequest> = command_channel.into();
    command_channel.readiness.insert(Ready::readable());

    if let Some(metrics) = worker_config.metrics.as_ref() {
        metrics::setup(
            &metrics.address,
            worker_id,
            metrics.tagged_metrics,
            metrics.prefix.clone(),
        );
    }

    let mut server = Server::new_from_config(
        command_channel,
        ScmSocket::new(scm_socket_fd),
        worker_config,
        config_state,
        true,
    );

    info!("starting event loop");
    server.run();
    info!("ending event loop");
    Ok(())
}

pub fn fork_main_into_worker(
    id: &str,
    config: &Config,
    executable_path: String,
    state: &ConfigState,
    listeners: Option<Listeners>,
) -> anyhow::Result<(pid_t, Channel<ProxyRequest, ProxyResponse>, ScmSocket)> {
    trace!("parent({})", unsafe { libc::getpid() });

    let mut state_file =
        tempfile().expect("could not create temporary file for configuration state");
    util::disable_close_on_exec(state_file.as_raw_fd())?;

    serde_json::to_writer(&mut state_file, state)
        .with_context(|| "could not write upgrade data to temporary file")?;
    state_file
        .seek(SeekFrom::Start(0))
        .with_context(|| "could not seek to beginning of file")?;

    let (command_socket_main, command_socket_worker) = UnixStream::pair()?;
    let (scm_socket_main, scm_socket_worker) = UnixStream::pair()?;

    let scm_main = ScmSocket::new(scm_socket_main.into_raw_fd());

    util::disable_close_on_exec(command_socket_worker.as_raw_fd())?;
    util::disable_close_on_exec(scm_socket_worker.as_raw_fd())?;

    let mut command_channel: Channel<Config, ProxyResponse> = Channel::new(
        command_socket_main,
        config.command_buffer_size,
        config.max_command_buffer_size,
    );
    command_channel.set_nonblocking(false);

    info!("{} launching worker", id);
    debug!("executable path is {}", executable_path);
    match unsafe { fork() } {
        Ok(ForkResult::Parent { child }) => {
            info!("{} worker launched: {}", id, child);
            command_channel.write_message(config);
            command_channel.set_nonblocking(true);

            if let Some(listeners) = listeners {
                info!("sending listeners to new worker: {:?}", listeners);
                let result = scm_main.send_listeners(&listeners);
                info!("sent listeners from main: {:?}", result);
                listeners.close();
            };
            util::disable_close_on_exec(scm_main.fd)?;

            Ok((child.into(), command_channel.into(), scm_main))
        }
        Ok(ForkResult::Child) => {
            trace!("child({}):\twill spawn a child", unsafe { libc::getpid() });
            Command::new(executable_path)
                .arg("worker")
                .arg("--id")
                .arg(id)
                .arg("--fd")
                .arg(command_socket_worker.as_raw_fd().to_string())
                .arg("--scm")
                .arg(scm_socket_worker.as_raw_fd().to_string())
                .arg("--configuration-state-fd")
                .arg(state_file.as_raw_fd().to_string())
                .arg("--command-buffer-size")
                .arg(config.command_buffer_size.to_string())
                .arg("--max-command-buffer-size")
                .arg(config.max_command_buffer_size.to_string())
                .exec();

            unreachable!();
        }
        Err(e) => {
            error!("Error during fork(): {}", e);
            Err(anyhow::Error::from(e))
        }
    }
}

#[cfg(target_os = "linux")]
pub unsafe fn get_executable_path() -> anyhow::Result<String> {
    use std::fs;

    let path = fs::read_link("/proc/self/exe").with_context(|| "/proc/self/exe doesn't exist")?;

    let mut path_str = match path.into_os_string().into_string() {
        Ok(s) => s,
        Err(_) => bail!("Failed to convert PathBuf to String"),
    };

    if path_str.ends_with(" (deleted)") {
        // The kernel appends " (deleted)" to the symlink when the original executable has been replaced
        let len = path_str.len();
        path_str.truncate(len - 10)
    }

    Ok(path_str)
}

#[cfg(target_os = "macos")]
extern "C" {
    pub fn _NSGetExecutablePath(buf: *mut c_char, size: *mut u32) -> i32;
}

#[cfg(target_os = "macos")]
pub unsafe fn get_executable_path() -> String {
    let capacity = PATH_MAX as usize;
    let mut temp: Vec<u8> = Vec::with_capacity(capacity);
    temp.extend(repeat(0).take(capacity));
    let pathbuf = CString::from_vec_unchecked(temp);
    let ptr = pathbuf.into_raw();

    let mut size = capacity as u32;
    if _NSGetExecutablePath(ptr, &mut size) == 0 {
        let mut temp2: Vec<u8> = Vec::with_capacity(capacity);
        temp2.extend(repeat(0).take(capacity));
        let pathbuf2 = CString::from_vec_unchecked(temp2);
        let ptr2 = pathbuf2.into_raw();

        if libc::realpath(ptr, ptr2) != null_mut() {
            let path = CString::from_raw(ptr2);
            path.to_str()
                .expect("failed to convert CString to String")
                .to_string()
        } else {
            panic!();
        }
    } else {
        panic!("buffer too small");
    }
}

#[cfg(target_os = "freebsd")]
pub unsafe fn get_executable_path() -> String {
    let mut capacity = PATH_MAX as usize;
    let mut path: Vec<u8> = Vec::with_capacity(capacity);
    path.extend(repeat(0).take(capacity));

    let mib: Vec<i32> = vec![CTL_KERN, KERN_PROC, KERN_PROC_PATHNAME, -1];
    let len = mib.len() * size_of::<i32>();
    let element_size = size_of::<i32>();

    let res = sysctl(
        mib.as_ptr(),
        (len / element_size) as u32,
        path.as_mut_ptr() as *mut c_void,
        &mut capacity,
        std::ptr::null() as *const c_void,
        0,
    );
    if res != 0 {
        panic!("Could not retrieve the path of the executable");
    }

    String::from_raw_parts(path.as_mut_ptr(), capacity - 1, path.len())
}
<|endoftext|>


<commit_before>use crate::data::messages::{BookUpdate, RecordUpdate, BookRecord};
use crate::data::trade::TradeBook;
use std::str::FromStr;
use super::Processor;
use std::sync::{Arc, Mutex};
use crate::error::PoloError;

#[derive(Clone)]
pub struct Accountant {
  tb: Arc<Mutex<TradeBook>>,
}

impl Accountant {
  pub fn new(tb: &Arc<Mutex<TradeBook>>) -> Accountant {
    Accountant {
      tb: tb.clone()
    }
  }
}

impl Processor for Accountant {
  fn process_message(&mut self, msg: String) -> Result<(),PoloError> {
    let err = |title| PoloError::wrong_data(format!("{} {:?}", title, msg));

    let update = BookUpdate::from_str(&msg)?;
    for rec in update.records {
      let mut tb = self.tb.lock().unwrap();
      match rec {
        RecordUpdate::Initial(book) => {
          tb.add_book(book, update.book_id);
        },
        RecordUpdate::SellTotal(BookRecord {rate, amount}) => {
          let book = tb.book_by_id(update.book_id)
            .ok_or_else(|| err("book not initialized"))?;
          book.update_sell_orders(rate, amount);
        },
        RecordUpdate::BuyTotal(BookRecord {rate, amount}) => {
          let book = tb.book_by_id(update.book_id)
            .ok_or_else(|| err("book not initialized"))?;
          book.update_buy_orders(rate, amount);
        },
        RecordUpdate::Sell(deal) => {
          let book = tb.book_by_id(update.book_id)
            .ok_or_else(|| err("book not initialized"))?;
          book.new_deal(deal.id, deal.rate, -deal.amount)?;
        },
        RecordUpdate::Buy(deal) => {
          let book = tb.book_by_id(update.book_id)
            .ok_or_else(|| err("book not initialized"))?;
          book.new_deal(deal.id, deal.rate, deal.amount)?;
        },
      }
    };
    Ok(())
  }
}

#[cfg(test)]
mod tests {
  use super::Accountant;
  use crate::actors::Processor;
  use crate::data::trade::TradeBook;
  use crate::data::messages::{BookUpdate, RecordUpdate};
  use std::str::FromStr;
  use std::sync::{Arc, Mutex};

  #[test]
  fn initial_order() {
    let tb = Arc::new(Mutex::new(TradeBook::new()));
    let mut accountant = Accountant::new(&tb);
    let order = String::from(r#"[189, 5130995, [["i", {"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13169621": 0.2331}]}]]]"#);

    accountant.process_message(order.clone()).unwrap();

    let mut tb_mut = tb.lock().unwrap();
    let actor_book = tb_mut.book_by_id(189).unwrap().book_ref();
    match BookUpdate::from_str(&order).unwrap().records[0] {
      RecordUpdate::Initial(ref book) => assert_eq!(*book, *actor_book),
      _ => panic!("BookUpdate::from_str were not able to parse RecordUpdate::Initial")
    }
  }
}<commit_msg>get rid of extra clone<commit_after>use crate::data::messages::{BookUpdate, RecordUpdate, BookRecord};
use crate::data::trade::TradeBook;
use std::str::FromStr;
use super::Processor;
use std::sync::{Arc, Mutex};
use crate::error::PoloError;

#[derive(Clone)]
pub struct Accountant {
  tb: Arc<Mutex<TradeBook>>,
}

impl Accountant {
  pub fn new(tb: Arc<Mutex<TradeBook>>) -> Accountant {
    Accountant { tb }
  }
}

impl Processor for Accountant {
  fn process_message(&mut self, msg: String) -> Result<(),PoloError> {
    let err = |title| PoloError::wrong_data(format!("{} {:?}", title, msg));

    let update = BookUpdate::from_str(&msg)?;
    for rec in update.records {
      let mut tb = self.tb.lock().unwrap();
      match rec {
        RecordUpdate::Initial(book) => {
          tb.add_book(book, update.book_id);
        },
        RecordUpdate::SellTotal(BookRecord {rate, amount}) => {
          let book = tb.book_by_id(update.book_id)
            .ok_or_else(|| err("book not initialized"))?;
          book.update_sell_orders(rate, amount);
        },
        RecordUpdate::BuyTotal(BookRecord {rate, amount}) => {
          let book = tb.book_by_id(update.book_id)
            .ok_or_else(|| err("book not initialized"))?;
          book.update_buy_orders(rate, amount);
        },
        RecordUpdate::Sell(deal) => {
          let book = tb.book_by_id(update.book_id)
            .ok_or_else(|| err("book not initialized"))?;
          book.new_deal(deal.id, deal.rate, -deal.amount)?;
        },
        RecordUpdate::Buy(deal) => {
          let book = tb.book_by_id(update.book_id)
            .ok_or_else(|| err("book not initialized"))?;
          book.new_deal(deal.id, deal.rate, deal.amount)?;
        },
      }
    };
    Ok(())
  }
}

#[cfg(test)]
mod tests {
  use super::Accountant;
  use crate::actors::Processor;
  use crate::data::trade::TradeBook;
  use crate::data::messages::{BookUpdate, RecordUpdate};
  use std::str::FromStr;
  use std::sync::{Arc, Mutex};

  #[test]
  fn initial_order() {
    let tb = Arc::new(Mutex::new(TradeBook::new()));
    let mut accountant = Accountant::new(&tb);
    let order = String::from(r#"[189, 5130995, [["i", {"currencyPair": "BTC_BCH", "orderBook": [{"0.13161901": 0.23709568, "0.13164313": "0.17328089"}, {"0.13169621": 0.2331}]}]]]"#);

    accountant.process_message(order.clone()).unwrap();

    let mut tb_mut = tb.lock().unwrap();
    let actor_book = tb_mut.book_by_id(189).unwrap().book_ref();
    match BookUpdate::from_str(&order).unwrap().records[0] {
      RecordUpdate::Initial(ref book) => assert_eq!(*book, *actor_book),
      _ => panic!("BookUpdate::from_str were not able to parse RecordUpdate::Initial")
    }
  }
}<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::collections::hash_map::{Entry, HashMap};
use std::collections::hash_set::HashSet;
use std::rc::{Rc, Weak};
use std::cmp;
use std::mem;

use js_types::js_type::{JsVar, JsType, JsPtrEnum};
use uuid::Uuid;

// Initial Arena size in bytes
const INITIAL_SIZE: usize = 1024;
// Minimum Arena capacity is at least 1 byte
const MIN_CAP: usize = 1;

pub type Alloc<T> = Rc<RefCell<T>>;

pub struct Scope {
    pub parent: Option<Weak<Scope>>,
    children: Vec<Rc<Scope>>,
    black_set: HashMap<Uuid, Alloc<JsVar>>,
    grey_set: HashMap<Uuid, Alloc<JsVar>>,
    white_set: HashMap<Uuid, Alloc<JsVar>>,
    get_roots: Box<Fn() -> HashSet<Uuid>>,
}

impl Scope {
    pub fn new<F>(get_roots: F) -> Scope
        where F: Fn() -> HashSet<Uuid> + 'static {
        Scope {
            parent: None,
            children: Vec::new(),
            black_set: HashMap::new(),
            grey_set: HashMap::new(),
            white_set: HashMap::new(),
            get_roots: Box::new(get_roots),
        }
    }

    pub fn as_child<F>(parent: Weak<Scope>, get_roots: F) -> Scope
        where F: Fn() -> HashSet<Uuid> + 'static {
        Scope {
            parent: Some(parent),
            children: Vec::new(),
            black_set: HashMap::new(),
            grey_set: HashMap::new(),
            white_set: HashMap::new(),
            get_roots: Box::new(get_roots),
        }
    }

    pub fn set_parent(&mut self, parent: &Rc<Scope>) {
        self.parent = Some(Rc::downgrade(&parent.clone()));
    }

    pub fn add_child(&mut self, child: Scope) -> &mut Rc<Scope> {
        self.children.push(Rc::new(child));
        let num_children = self.children.len();
        &mut self.children[num_children - 1]
    }

    pub fn alloc(&mut self, var: JsVar) -> Uuid {
        let uuid = var.uuid;
        self.white_set.insert(uuid, Rc::new(RefCell::new(var)));
        uuid
    }

    pub fn dealloc(&mut self, uuid: &Uuid) -> bool {
        if let Some(_) = self.white_set.remove(uuid) { true } else { false }
    }

    pub fn get_var_copy(&self, uuid: &Uuid) -> Option<JsVar> {
        self.find_id(uuid).map(|var| var.borrow().clone())
    }

    pub fn update_var(&mut self, var: JsVar) -> bool {
        if let Entry::Occupied(mut view) = self.find_id_mut(&var.uuid) {
            let inner = view.get_mut();
            *inner.borrow_mut() = var;
            true
        } else { false }
    }

    /// TODO Compute the roots of the current scope-- any variable that is
    /// directly referenced or declared within the scope. This might just be the
    /// key set of the uuid map(?) Not necessarily, I think. What if you do
    /// something like this:
    /// var x = {}
    /// var y = { 1: x }
    /// y = x
    /// y would be a root by the definition above, but is no longer reachable at
    /// the end of the scope because it now aliases x. A better definition would
    /// be "Any variable that is declared or referenced directly, but a direct
    /// reference (variable usage) supercedes a declaration." The above example
    /// demonstrates why this is necessary.
    /// This should come from the interpreter, so I shouldn't actually have to
    /// care about getting the root set myself.

    /// Roots always get marked as Black, since they're always reachable from
    /// the current scope. NB that this assumes all root references are actually
    /// valid reference types, i.e. they're not numbers, etc.
    pub fn mark_roots(&mut self) {
        let marks = (self.get_roots)();
        for mark in marks.iter() {
            if let Some(var) = self.white_set.remove(mark) {
                let uuid = var.borrow().uuid;
                // Get all child references
                let child_ids = Scope::get_var_children(&var);
                // Mark current ref as black
                self.black_set.insert(uuid, var);
                // Mark child references as grey
                self.grey_children(child_ids);
            }
        }
    }

    pub fn mark_phase(&mut self) {
        // Mark any grey object as black, and mark all white objs it refs as grey
        let mut new_grey_set = HashMap::new();
        for (uuid, var) in self.grey_set.drain() {
            let child_ids = Scope::get_var_children(&var);
            self.black_set.insert(uuid, var);
            for child_id in child_ids {
                if let Some(var) = self.white_set.remove(&child_id) {
                    new_grey_set.insert(child_id, var);
                }
            }
        }
        self.grey_set = new_grey_set;
    }

    pub fn sweep_phase(&mut self) {
        self.white_set = HashMap::new();
    }

    pub fn find_id(&self, uuid: &Uuid) -> Option<&Alloc<JsVar>> {
        self.white_set.get(uuid).or_else(||
            self.grey_set.get(uuid).or_else(||
                self.black_set.get(uuid)))
    }

    fn find_id_mut(&mut self, uuid: &Uuid) -> Entry<Uuid, Alloc<JsVar>> {
        if let e @ Entry::Occupied(_) = self.white_set.entry(*uuid) {
            e
        } else if let e @ Entry::Occupied(_) = self.grey_set.entry(*uuid) {
            e
        } else { self.black_set.entry(*uuid) }
    }

    fn grey_children(&mut self, child_ids: HashSet<Uuid>) {
        for child_id in child_ids.iter() {
            if let Some(var) = self.white_set.remove(child_id) {
                self.grey_set.insert(*child_id, var);
            }
        }
    }

    fn get_var_children(var: &Alloc<JsVar>) -> HashSet<Uuid> {
        if let JsType::JsPtr(ref ptr) = (*var.borrow()).t {
            match ptr {
                &JsPtrEnum::JsObj(ref obj) => obj.get_children(),
                _ => HashSet::new(),
            }
        } else { HashSet::new() }
    }

}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::hash_set::HashSet;
    use std::rc::{Rc, Weak};
    use js_types::js_type::{JsVar, JsType, JsPtrEnum, JsKey, JsKeyEnum};
    use js_types::js_obj::{JsObjStruct};
    use uuid::Uuid;

    fn dummy_get_roots() -> HashSet<Uuid> {
        HashSet::new()
    }

    fn make_num(i: f64) -> JsVar {
        JsVar::new(JsType::JsNum(i))
    }

    fn make_obj(kvs: Vec<(JsKey, Alloc<JsVar>)>) -> JsVar {
        JsVar::new(JsType::JsPtr(JsPtrEnum::JsObj(JsObjStruct::new(None, "test", kvs))))
    }

    #[test]
    fn test_new_scope() {
        let mut test_scope = Scope::new(dummy_get_roots);
        assert!(test_scope.parent.is_none());
        assert!(test_scope.black_set.is_empty());
        assert!(test_scope.grey_set.is_empty());
        assert!(test_scope.white_set.is_empty());
        assert_eq!(test_scope.children.len(), 0);
    }

    #[test]
    fn test_as_child_scope() {
        let parent_scope = Rc::new(Scope::new(dummy_get_roots));
        let mut test_scope = Scope::as_child(Rc::downgrade(&parent_scope.clone()), dummy_get_roots);

        assert!(test_scope.parent.is_some());
        assert!(test_scope.black_set.is_empty());
        assert!(test_scope.grey_set.is_empty());
        assert!(test_scope.white_set.is_empty());
        assert_eq!(test_scope.children.len(), 0);
    }

    #[test]
    fn test_set_parent() {
        let parent_scope = Rc::new(Scope::new(dummy_get_roots));
        let mut test_scope = Scope::new(dummy_get_roots);
        assert!(test_scope.parent.is_none());
        test_scope.set_parent(&parent_scope);
        assert!(test_scope.parent.is_some());
    }

    #[test]
    fn test_add_child() {
        let mut test_scope = Scope::new(dummy_get_roots);
        let child_scope1 = Scope::new(dummy_get_roots);
        let child_scope2 = Scope::new(dummy_get_roots);
        assert_eq!(test_scope.children.len(), 0);
        test_scope.add_child(child_scope1);
        assert_eq!(test_scope.children.len(), 1);
        test_scope.add_child(child_scope2);
        assert_eq!(test_scope.children.len(), 2);
    }

    #[test]
    fn test_alloc() {
        let mut test_scope = Scope::new(dummy_get_roots);
        let test_var = make_num(1.0);
        let test_uuid = test_var.uuid.clone();
        let uuid = test_scope.alloc(test_var);
        assert_eq!(test_uuid, uuid);
        assert!(test_scope.white_set.contains_key(&uuid));
        assert_eq!(test_scope.white_set.len(), 1);
        assert_eq!(test_scope.grey_set.len(), 0);
        assert_eq!(test_scope.black_set.len(), 0);
        let test_var2 = make_num(2.0);
        let uuid = test_scope.alloc(test_var2);
        assert!(test_scope.white_set.contains_key(&uuid));
        assert_eq!(test_scope.white_set.len(), 2);
        assert_eq!(test_scope.grey_set.len(), 0);
        assert_eq!(test_scope.black_set.len(), 0);
    }

    #[test]
    fn test_dealloc() {
        let mut test_scope = Scope::new(dummy_get_roots);
        let test_var = make_num(1.0);
        let uuid = test_scope.alloc(test_var);
        let bad_uuid = Uuid::new_v4();
        assert!(test_scope.dealloc(&uuid));
        assert_eq!(test_scope.white_set.len(), 0);
        assert_eq!(test_scope.grey_set.len(), 0);
        assert_eq!(test_scope.black_set.len(), 0);
        assert!(!test_scope.dealloc(&bad_uuid));
    }

    #[test]
    fn test_get_var_copy() {
        let mut test_scope = Scope::new(dummy_get_roots);
        let test_var = make_num(1.0);
        let uuid = test_scope.alloc(test_var);
        let bad_uuid = Uuid::new_v4();
        let var_copy = test_scope.get_var_copy(&uuid);
        assert!(var_copy.is_some());
        let var = var_copy.unwrap();
        assert_eq!(var.uuid, uuid);
        let bad_copy = test_scope.get_var_copy(&bad_uuid);
        assert!(bad_copy.is_none());
    }

    #[test]
    fn test_update_var() {
        let mut test_scope = Scope::new(dummy_get_roots);
        let test_var = make_num(1.0);
        let uuid = test_scope.alloc(test_var);
        let mut update = test_scope.get_var_copy(&uuid).unwrap();
        update = make_num(2.0);
        update.uuid = uuid;
        assert!(test_scope.update_var(update));
        let update = test_scope.get_var_copy(&uuid).unwrap();
        match update {
            JsVar{ t: JsType::JsNum(i), ..} => assert_eq!(i, 2.0),
            _ => ()
        }
        test_scope.dealloc(&uuid);
        assert!(!test_scope.update_var(update));
    }

    #[test]
    fn test_mark_roots() {
        let test_uuid = Uuid::new_v4();
        let mut test_var = make_num(1.0);
        let mut test_var2 = make_num(2.0);
        test_var.uuid = test_uuid;
        let simple_get_roots = move || { let mut set = HashSet::new();
                                    set.insert(test_uuid);
                                    set };
        let mut test_scope = Scope::new(simple_get_roots);
        let test_uuid = test_scope.alloc(test_var);
        let test_uuid2 = test_scope.alloc(test_var2);
        test_scope.mark_roots();
        assert_eq!(test_scope.black_set.len(), 1);
        assert!(test_scope.black_set.get(&test_uuid).is_some());
        assert!(test_scope.black_set.get(&test_uuid2).is_none());
    }

    #[test]
    fn test_mark_phase() {
        let test_uuid = Uuid::new_v4();
        let simple_get_roots = move || { let mut set = HashSet::new();
                                    set.insert(test_uuid);
                                    set };
        let mut test_scope = Scope::new(simple_get_roots);

        let mut test_var2 = make_num(2.0);
        let test_uuid2 = test_scope.alloc(test_var2);
        let test_var2_alloc = test_scope.find_id(&test_uuid2).unwrap().clone();

        let test_key = JsKey::new(JsKeyEnum::JsNum(1.0));
        let mut test_var = make_obj(vec![(test_key, test_var2_alloc)]);
        test_var.uuid = test_uuid;
        let test_uuid = test_scope.alloc(test_var);

        // Mark roots first...
        test_scope.mark_roots();
        assert!(test_scope.black_set.get(&test_uuid).is_some());
        assert_eq!(test_scope.black_set.len(), 1);
        assert_eq!(test_scope.grey_set.len(), 1);

        // ...then mark everything else
        test_scope.mark_phase();

        assert_eq!(test_scope.grey_set.len(), 0);
        assert!(test_scope.black_set.get(&test_uuid).is_some());
        assert!(test_scope.black_set.get(&test_uuid2).is_some());
    }
}
<commit_msg>Make Scope children public<commit_after>use std::cell::RefCell;
use std::collections::hash_map::{Entry, HashMap};
use std::collections::hash_set::HashSet;
use std::rc::{Rc, Weak};
use std::cmp;
use std::mem;

use js_types::js_type::{JsVar, JsType, JsPtrEnum};
use uuid::Uuid;

// Initial Arena size in bytes
const INITIAL_SIZE: usize = 1024;
// Minimum Arena capacity is at least 1 byte
const MIN_CAP: usize = 1;

pub type Alloc<T> = Rc<RefCell<T>>;

pub struct Scope {
    pub parent: Option<Weak<Scope>>,
    pub children: Vec<Rc<Scope>>,
    black_set: HashMap<Uuid, Alloc<JsVar>>,
    grey_set: HashMap<Uuid, Alloc<JsVar>>,
    white_set: HashMap<Uuid, Alloc<JsVar>>,
    get_roots: Box<Fn() -> HashSet<Uuid>>,
}

impl Scope {
    pub fn new<F>(get_roots: F) -> Scope
        where F: Fn() -> HashSet<Uuid> + 'static {
        Scope {
            parent: None,
            children: Vec::new(),
            black_set: HashMap::new(),
            grey_set: HashMap::new(),
            white_set: HashMap::new(),
            get_roots: Box::new(get_roots),
        }
    }

    pub fn as_child<F>(parent: Weak<Scope>, get_roots: F) -> Scope
        where F: Fn() -> HashSet<Uuid> + 'static {
        Scope {
            parent: Some(parent),
            children: Vec::new(),
            black_set: HashMap::new(),
            grey_set: HashMap::new(),
            white_set: HashMap::new(),
            get_roots: Box::new(get_roots),
        }
    }

    pub fn set_parent(&mut self, parent: &Rc<Scope>) {
        self.parent = Some(Rc::downgrade(&parent.clone()));
    }

    pub fn add_child(&mut self, child: Scope) -> &mut Rc<Scope> {
        self.children.push(Rc::new(child));
        let num_children = self.children.len();
        &mut self.children[num_children - 1]
    }

    pub fn alloc(&mut self, var: JsVar) -> Uuid {
        let uuid = var.uuid;
        self.white_set.insert(uuid, Rc::new(RefCell::new(var)));
        uuid
    }

    pub fn dealloc(&mut self, uuid: &Uuid) -> bool {
        if let Some(_) = self.white_set.remove(uuid) { true } else { false }
    }

    pub fn get_var_copy(&self, uuid: &Uuid) -> Option<JsVar> {
        self.find_id(uuid).map(|var| var.borrow().clone())
    }

    pub fn update_var(&mut self, var: JsVar) -> bool {
        if let Entry::Occupied(mut view) = self.find_id_mut(&var.uuid) {
            let inner = view.get_mut();
            *inner.borrow_mut() = var;
            true
        } else { false }
    }

    /// TODO Compute the roots of the current scope-- any variable that is
    /// directly referenced or declared within the scope. This might just be the
    /// key set of the uuid map(?) Not necessarily, I think. What if you do
    /// something like this:
    /// var x = {}
    /// var y = { 1: x }
    /// y = x
    /// y would be a root by the definition above, but is no longer reachable at
    /// the end of the scope because it now aliases x. A better definition would
    /// be "Any variable that is declared or referenced directly, but a direct
    /// reference (variable usage) supercedes a declaration." The above example
    /// demonstrates why this is necessary.
    /// This should come from the interpreter, so I shouldn't actually have to
    /// care about getting the root set myself.

    /// Roots always get marked as Black, since they're always reachable from
    /// the current scope. NB that this assumes all root references are actually
    /// valid reference types, i.e. they're not numbers, etc.
    pub fn mark_roots(&mut self) {
        let marks = (self.get_roots)();
        for mark in marks.iter() {
            if let Some(var) = self.white_set.remove(mark) {
                let uuid = var.borrow().uuid;
                // Get all child references
                let child_ids = Scope::get_var_children(&var);
                // Mark current ref as black
                self.black_set.insert(uuid, var);
                // Mark child references as grey
                self.grey_children(child_ids);
            }
        }
    }

    pub fn mark_phase(&mut self) {
        // Mark any grey object as black, and mark all white objs it refs as grey
        let mut new_grey_set = HashMap::new();
        for (uuid, var) in self.grey_set.drain() {
            let child_ids = Scope::get_var_children(&var);
            self.black_set.insert(uuid, var);
            for child_id in child_ids {
                if let Some(var) = self.white_set.remove(&child_id) {
                    new_grey_set.insert(child_id, var);
                }
            }
        }
        self.grey_set = new_grey_set;
    }

    pub fn sweep_phase(&mut self) {
        self.white_set = HashMap::new();
    }

    pub fn find_id(&self, uuid: &Uuid) -> Option<&Alloc<JsVar>> {
        self.white_set.get(uuid).or_else(||
            self.grey_set.get(uuid).or_else(||
                self.black_set.get(uuid)))
    }

    fn find_id_mut(&mut self, uuid: &Uuid) -> Entry<Uuid, Alloc<JsVar>> {
        if let e @ Entry::Occupied(_) = self.white_set.entry(*uuid) {
            e
        } else if let e @ Entry::Occupied(_) = self.grey_set.entry(*uuid) {
            e
        } else { self.black_set.entry(*uuid) }
    }

    fn grey_children(&mut self, child_ids: HashSet<Uuid>) {
        for child_id in child_ids.iter() {
            if let Some(var) = self.white_set.remove(child_id) {
                self.grey_set.insert(*child_id, var);
            }
        }
    }

    fn get_var_children(var: &Alloc<JsVar>) -> HashSet<Uuid> {
        if let JsType::JsPtr(ref ptr) = (*var.borrow()).t {
            match ptr {
                &JsPtrEnum::JsObj(ref obj) => obj.get_children(),
                _ => HashSet::new(),
            }
        } else { HashSet::new() }
    }

}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::hash_set::HashSet;
    use std::rc::{Rc, Weak};
    use js_types::js_type::{JsVar, JsType, JsPtrEnum, JsKey, JsKeyEnum};
    use js_types::js_obj::{JsObjStruct};
    use uuid::Uuid;

    fn dummy_get_roots() -> HashSet<Uuid> {
        HashSet::new()
    }

    fn make_num(i: f64) -> JsVar {
        JsVar::new(JsType::JsNum(i))
    }

    fn make_obj(kvs: Vec<(JsKey, Alloc<JsVar>)>) -> JsVar {
        JsVar::new(JsType::JsPtr(JsPtrEnum::JsObj(JsObjStruct::new(None, "test", kvs))))
    }

    #[test]
    fn test_new_scope() {
        let mut test_scope = Scope::new(dummy_get_roots);
        assert!(test_scope.parent.is_none());
        assert!(test_scope.black_set.is_empty());
        assert!(test_scope.grey_set.is_empty());
        assert!(test_scope.white_set.is_empty());
        assert_eq!(test_scope.children.len(), 0);
    }

    #[test]
    fn test_as_child_scope() {
        let parent_scope = Rc::new(Scope::new(dummy_get_roots));
        let mut test_scope = Scope::as_child(Rc::downgrade(&parent_scope.clone()), dummy_get_roots);

        assert!(test_scope.parent.is_some());
        assert!(test_scope.black_set.is_empty());
        assert!(test_scope.grey_set.is_empty());
        assert!(test_scope.white_set.is_empty());
        assert_eq!(test_scope.children.len(), 0);
    }

    #[test]
    fn test_set_parent() {
        let parent_scope = Rc::new(Scope::new(dummy_get_roots));
        let mut test_scope = Scope::new(dummy_get_roots);
        assert!(test_scope.parent.is_none());
        test_scope.set_parent(&parent_scope);
        assert!(test_scope.parent.is_some());
    }

    #[test]
    fn test_add_child() {
        let mut test_scope = Scope::new(dummy_get_roots);
        let child_scope1 = Scope::new(dummy_get_roots);
        let child_scope2 = Scope::new(dummy_get_roots);
        assert_eq!(test_scope.children.len(), 0);
        test_scope.add_child(child_scope1);
        assert_eq!(test_scope.children.len(), 1);
        test_scope.add_child(child_scope2);
        assert_eq!(test_scope.children.len(), 2);
    }

    #[test]
    fn test_alloc() {
        let mut test_scope = Scope::new(dummy_get_roots);
        let test_var = make_num(1.0);
        let test_uuid = test_var.uuid.clone();
        let uuid = test_scope.alloc(test_var);
        assert_eq!(test_uuid, uuid);
        assert!(test_scope.white_set.contains_key(&uuid));
        assert_eq!(test_scope.white_set.len(), 1);
        assert_eq!(test_scope.grey_set.len(), 0);
        assert_eq!(test_scope.black_set.len(), 0);
        let test_var2 = make_num(2.0);
        let uuid = test_scope.alloc(test_var2);
        assert!(test_scope.white_set.contains_key(&uuid));
        assert_eq!(test_scope.white_set.len(), 2);
        assert_eq!(test_scope.grey_set.len(), 0);
        assert_eq!(test_scope.black_set.len(), 0);
    }

    #[test]
    fn test_dealloc() {
        let mut test_scope = Scope::new(dummy_get_roots);
        let test_var = make_num(1.0);
        let uuid = test_scope.alloc(test_var);
        let bad_uuid = Uuid::new_v4();
        assert!(test_scope.dealloc(&uuid));
        assert_eq!(test_scope.white_set.len(), 0);
        assert_eq!(test_scope.grey_set.len(), 0);
        assert_eq!(test_scope.black_set.len(), 0);
        assert!(!test_scope.dealloc(&bad_uuid));
    }

    #[test]
    fn test_get_var_copy() {
        let mut test_scope = Scope::new(dummy_get_roots);
        let test_var = make_num(1.0);
        let uuid = test_scope.alloc(test_var);
        let bad_uuid = Uuid::new_v4();
        let var_copy = test_scope.get_var_copy(&uuid);
        assert!(var_copy.is_some());
        let var = var_copy.unwrap();
        assert_eq!(var.uuid, uuid);
        let bad_copy = test_scope.get_var_copy(&bad_uuid);
        assert!(bad_copy.is_none());
    }

    #[test]
    fn test_update_var() {
        let mut test_scope = Scope::new(dummy_get_roots);
        let test_var = make_num(1.0);
        let uuid = test_scope.alloc(test_var);
        let mut update = test_scope.get_var_copy(&uuid).unwrap();
        update = make_num(2.0);
        update.uuid = uuid;
        assert!(test_scope.update_var(update));
        let update = test_scope.get_var_copy(&uuid).unwrap();
        match update {
            JsVar{ t: JsType::JsNum(i), ..} => assert_eq!(i, 2.0),
            _ => ()
        }
        test_scope.dealloc(&uuid);
        assert!(!test_scope.update_var(update));
    }

    #[test]
    fn test_mark_roots() {
        let test_uuid = Uuid::new_v4();
        let mut test_var = make_num(1.0);
        let mut test_var2 = make_num(2.0);
        test_var.uuid = test_uuid;
        let simple_get_roots = move || { let mut set = HashSet::new();
                                    set.insert(test_uuid);
                                    set };
        let mut test_scope = Scope::new(simple_get_roots);
        let test_uuid = test_scope.alloc(test_var);
        let test_uuid2 = test_scope.alloc(test_var2);
        test_scope.mark_roots();
        assert_eq!(test_scope.black_set.len(), 1);
        assert!(test_scope.black_set.get(&test_uuid).is_some());
        assert!(test_scope.black_set.get(&test_uuid2).is_none());
    }

    #[test]
    fn test_mark_phase() {
        let test_uuid = Uuid::new_v4();
        let simple_get_roots = move || { let mut set = HashSet::new();
                                    set.insert(test_uuid);
                                    set };
        let mut test_scope = Scope::new(simple_get_roots);

        let mut test_var2 = make_num(2.0);
        let test_uuid2 = test_scope.alloc(test_var2);
        let test_var2_alloc = test_scope.find_id(&test_uuid2).unwrap().clone();

        let test_key = JsKey::new(JsKeyEnum::JsNum(1.0));
        let mut test_var = make_obj(vec![(test_key, test_var2_alloc)]);
        test_var.uuid = test_uuid;
        let test_uuid = test_scope.alloc(test_var);

        // Mark roots first...
        test_scope.mark_roots();
        assert!(test_scope.black_set.get(&test_uuid).is_some());
        assert_eq!(test_scope.black_set.len(), 1);
        assert_eq!(test_scope.grey_set.len(), 1);

        // ...then mark everything else
        test_scope.mark_phase();

        assert_eq!(test_scope.grey_set.len(), 0);
        assert!(test_scope.black_set.get(&test_uuid).is_some());
        assert!(test_scope.black_set.get(&test_uuid2).is_some());
    }
}
<|endoftext|>


<commit_before>#[naked]
pub unsafe extern fn syscall() {
    #[inline(never)]
    unsafe fn inner() {
        extern {
            fn syscall(a: usize, b: usize, c: usize, d: usize, e: usize, f: usize, stack: usize) -> usize;
        }

        let mut a;
        {
            let b;
            let c;
            let d;
            let e;
            let f;
            let stack;
            asm!("" : "={rax}"(a), "={rbx}"(b), "={rcx}"(c), "={rdx}"(d), "={rsi}"(e), "={rdi}"(f), "={rbp}"(stack)
                : : : "intel", "volatile");

            a = syscall(a, b, c, d, e, f, stack);
        }

        asm!("" : : "{rax}"(a) : : "intel", "volatile");
    }

    asm!("push r15
        push fs
        mov r15, 0x18
        mov fs, r15"
        : : : : "intel", "volatile");

    inner();

    // Interrupt return
    asm!("pop fs
        pop r15
        iretq"
        : : : : "intel", "volatile");
}

#[naked]
pub unsafe extern fn clone_ret() -> usize {
    asm!("pop rbp"
        : : : : "intel", "volatile");
        0
}
<commit_msg>Fix syscall ABI<commit_after>#[naked]
pub unsafe extern fn syscall() {
    #[inline(never)]
    unsafe fn inner() {
        extern {
            fn syscall(a: usize, b: usize, c: usize, d: usize, e: usize, f: usize, stack: usize) -> usize;
        }

        let mut a;
        {
            let b;
            let c;
            let d;
            let e;
            let f;
            let stack;
            asm!("" : "={rax}"(a), "={rbx}"(b), "={rcx}"(c), "={rdx}"(d), "={rsi}"(e), "={rdi}"(f), "={rbp}"(stack)
                : : : "intel", "volatile");

            a = syscall(a, b, c, d, e, f, stack);
        }

        asm!("" : : "{rax}"(a) : : "intel", "volatile");
    }

    // Push scratch registers, minus rax for the return value
    asm!("push rcx
        push rdx
        push rdi
        push rsi
        push r8
        push r9
        push r10
        push r11
        push fs
        mov r11, 0x18
        mov fs, r11"
        : : : : "intel", "volatile");

    inner();

    // Interrupt return
    asm!("pop fs
        pop r11
        pop r10
        pop r9
        pop r8
        pop rsi
        pop rdi
        pop rdx
        pop rcx
        iretq"
        : : : : "intel", "volatile");
}

#[naked]
pub unsafe extern fn clone_ret() -> usize {
    asm!("pop rbp"
        : : : : "intel", "volatile");
        0
}
<|endoftext|>


<commit_before>use artano::{self, Canvas, Typeface};
use config::{Options, OutputFormat};
use error::Error;
use std::path::Path;
use std::result;

type Result<T> = result::Result<T, Error>;

pub struct App;

impl App {
    pub fn run(&self, options: &Options) -> Result<()> {
        let font = build_font(&options.font_path)?;
        let mut canvas = options
            .base_image
            .get()
            .map_err(|e| Error::not_found("Base image not found", e))
            .and_then(|buf| {
                Canvas::read_from_buffer(&buf).map_err(Error::bad_image)
            })?;

        for annotation in &options.annotations {
            canvas.add_annotation(annotation, &font, options.scale_mult);
        }

        canvas.render();
        save_pixels(&options.output_path, &canvas, options.output_format)
    }
}

fn build_font(path: &Path) -> Result<Typeface> {
    use std::fs::File;
    use std::io::BufReader;

    let data = File::open(path).map_err(
        |e| Error::not_found("Font not found", e),
    )?;

    artano::load_typeface(&mut BufReader::new(data)).map_err(
        |e| {
            Error::io("Unable to read font", e)
        },
    )
}

fn save_pixels<P: AsRef<Path>>(path: P, canvas: &Canvas, format: OutputFormat) -> Result<()> {
    use std::fs::OpenOptions;

    let mut out = OpenOptions::new()
        .write(true)
        .create(true)
        .truncate(true)
        .open(path.as_ref())
        .map_err(|e| Error::io("Unable to write to output", e))?;

    let result = match format {
        OutputFormat::Png => canvas.save_png(&mut out),
        OutputFormat::Jpg => canvas.save_jpg(&mut out),
    };

    result.map_err(|e| Error::io("Unable to save image to output", e))
}
<commit_msg>Piddling<commit_after>use artano::{self, Canvas, Typeface};
use config::{Options, OutputFormat};
use error::Error;
use std::path::Path;
use std::result;

type Result<T> = result::Result<T, Error>;

pub struct App;

impl App {
    pub fn run(&self, options: &Options) -> Result<()> {
        let font = build_font(&options.font_path)?;
        let mut canvas = options
            .base_image
            .get()
            .map_err(|e| Error::not_found("Base image not found", e))
            .and_then(|buf| {
                Canvas::read_from_buffer(&buf).map_err(Error::bad_image)
            })?;

        for annotation in &options.annotations {
            canvas.add_annotation(annotation, &font, options.scale_mult);
        }

        canvas.render();
        save_pixels(&options.output_path, &canvas, options.output_format)
    }
}

fn build_font(path: &Path) -> Result<Typeface> {
    use std::fs::File;
    use std::io::BufReader;

    let data = File::open(path)
        .map(BufReader::new)
        .map_err(|e| Error::not_found("Font not found", e))?;

    artano::load_typeface(data)
        .map_err(|e| Error::io("Unable to read font", e))
}

fn save_pixels<P: AsRef<Path>>(path: P, canvas: &Canvas, format: OutputFormat) -> Result<()> {
    use std::fs::OpenOptions;

    let mut out = OpenOptions::new()
        .write(true)
        .create(true)
        .truncate(true)
        .open(path.as_ref())
        .map_err(|e| Error::io("Unable to write to output", e))?;

    let result = match format {
        OutputFormat::Png => canvas.save_png(&mut out),
        OutputFormat::Jpg => canvas.save_jpg(&mut out),
    };

    result.map_err(|e| Error::io("Unable to save image to output", e))
}
<|endoftext|>


<commit_before>pub mod download;
pub mod extract;<commit_msg>Add verify to mod<commit_after>pub mod download;
pub mod extract;
pub mod verify;<|endoftext|>