<commit_before>#![cfg_attr(thiserror_nightly_testing, feature(backtrace))]

use thiserror::Error;

#[derive(Error, Debug)]
#[error("...")]
pub struct Inner;

#[cfg(thiserror_nightly_testing)]
pub mod structs {
    use super::Inner;
    use std::backtrace::Backtrace;
    use std::error::Error;
    use std::sync::Arc;
    use thiserror::Error;

    #[derive(Error, Debug)]
    #[error("...")]
    pub struct PlainBacktrace {
        backtrace: Backtrace,
    }

    #[derive(Error, Debug)]
    #[error("...")]
    pub struct ExplicitBacktrace {
        #[backtrace]
        backtrace: Backtrace,
    }

    #[derive(Error, Debug)]
    #[error("...")]
    pub struct OptBacktrace {
        #[backtrace]
        backtrace: Option<Backtrace>,
    }

    #[derive(Error, Debug)]
    #[error("...")]
    pub struct ArcBacktrace {
        #[backtrace]
        backtrace: Arc<Backtrace>,
    }

    #[derive(Error, Debug)]
    #[error("...")]
    pub struct BacktraceFrom {
        #[from]
        source: Inner,
        #[backtrace]
        backtrace: Backtrace,
    }

    #[derive(Error, Debug)]
    #[error("...")]
    pub struct OptBacktraceFrom {
        #[from]
        source: Inner,
        #[backtrace]
        backtrace: Option<Backtrace>,
    }

    #[test]
    fn test_backtrace() {
        let error = PlainBacktrace {
            backtrace: Backtrace::capture(),
        };
        assert!(error.backtrace().is_some());

        let error = ExplicitBacktrace {
            backtrace: Backtrace::capture(),
        };
        assert!(error.backtrace().is_some());

        let error = OptBacktrace {
            backtrace: Some(Backtrace::capture()),
        };
        assert!(error.backtrace().is_some());

        let error = ArcBacktrace {
            backtrace: Arc::new(Backtrace::capture()),
        };
        assert!(error.backtrace().is_some());

        let error = BacktraceFrom::from(Inner);
        assert!(error.backtrace().is_some());

        let error = OptBacktraceFrom::from(Inner);
        assert!(error.backtrace().is_some());
    }
}

#[cfg(thiserror_nightly_testing)]
pub mod enums {
    use super::Inner;
    use std::backtrace::Backtrace;
    use std::error::Error;
    use std::sync::Arc;
    use thiserror::Error;

    #[derive(Error, Debug)]
    pub enum PlainBacktrace {
        #[error("...")]
        Test { backtrace: Backtrace },
    }

    #[derive(Error, Debug)]
    pub enum ExplicitBacktrace {
        #[error("...")]
        Test {
            #[backtrace]
            backtrace: Backtrace,
        },
    }

    #[derive(Error, Debug)]
    pub enum OptBacktrace {
        #[error("...")]
        Test {
            #[backtrace]
            backtrace: Option<Backtrace>,
        },
    }

    #[derive(Error, Debug)]
    pub enum ArcBacktrace {
        #[error("...")]
        Test {
            #[backtrace]
            backtrace: Arc<Backtrace>,
        },
    }

    #[derive(Error, Debug)]
    pub enum BacktraceFrom {
        #[error("...")]
        Test {
            #[from]
            source: Inner,
            #[backtrace]
            backtrace: Backtrace,
        },
    }

    #[derive(Error, Debug)]
    pub enum OptBacktraceFrom {
        #[error("...")]
        Test {
            #[from]
            source: Inner,
            #[backtrace]
            backtrace: Option<Backtrace>,
        },
    }

    #[test]
    fn test_backtrace() {
        let error = PlainBacktrace::Test {
            backtrace: Backtrace::capture(),
        };
        assert!(error.backtrace().is_some());

        let error = ExplicitBacktrace::Test {
            backtrace: Backtrace::capture(),
        };
        assert!(error.backtrace().is_some());

        let error = OptBacktrace::Test {
            backtrace: Some(Backtrace::capture()),
        };
        assert!(error.backtrace().is_some());

        let error = ArcBacktrace::Test {
            backtrace: Arc::new(Backtrace::capture()),
        };
        assert!(error.backtrace().is_some());

        let error = BacktraceFrom::from(Inner);
        assert!(error.backtrace().is_some());

        let error = OptBacktraceFrom::from(Inner);
        assert!(error.backtrace().is_some());
    }
}

#[test]
#[cfg_attr(not(thiserror_nightly_testing), ignore)]
fn test_backtrace() {}
<commit_msg>Add test of capturing Arc<Backtrace><commit_after>#![cfg_attr(thiserror_nightly_testing, feature(backtrace))]

use thiserror::Error;

#[derive(Error, Debug)]
#[error("...")]
pub struct Inner;

#[cfg(thiserror_nightly_testing)]
pub mod structs {
    use super::Inner;
    use std::backtrace::Backtrace;
    use std::error::Error;
    use std::sync::Arc;
    use thiserror::Error;

    #[derive(Error, Debug)]
    #[error("...")]
    pub struct PlainBacktrace {
        backtrace: Backtrace,
    }

    #[derive(Error, Debug)]
    #[error("...")]
    pub struct ExplicitBacktrace {
        #[backtrace]
        backtrace: Backtrace,
    }

    #[derive(Error, Debug)]
    #[error("...")]
    pub struct OptBacktrace {
        #[backtrace]
        backtrace: Option<Backtrace>,
    }

    #[derive(Error, Debug)]
    #[error("...")]
    pub struct ArcBacktrace {
        #[backtrace]
        backtrace: Arc<Backtrace>,
    }

    #[derive(Error, Debug)]
    #[error("...")]
    pub struct BacktraceFrom {
        #[from]
        source: Inner,
        #[backtrace]
        backtrace: Backtrace,
    }

    #[derive(Error, Debug)]
    #[error("...")]
    pub struct OptBacktraceFrom {
        #[from]
        source: Inner,
        #[backtrace]
        backtrace: Option<Backtrace>,
    }

    #[derive(Error, Debug)]
    #[error("...")]
    pub struct ArcBacktraceFrom {
        #[from]
        source: Inner,
        #[backtrace]
        backtrace: Arc<Backtrace>,
    }

    #[test]
    fn test_backtrace() {
        let error = PlainBacktrace {
            backtrace: Backtrace::capture(),
        };
        assert!(error.backtrace().is_some());

        let error = ExplicitBacktrace {
            backtrace: Backtrace::capture(),
        };
        assert!(error.backtrace().is_some());

        let error = OptBacktrace {
            backtrace: Some(Backtrace::capture()),
        };
        assert!(error.backtrace().is_some());

        let error = ArcBacktrace {
            backtrace: Arc::new(Backtrace::capture()),
        };
        assert!(error.backtrace().is_some());

        let error = BacktraceFrom::from(Inner);
        assert!(error.backtrace().is_some());

        let error = OptBacktraceFrom::from(Inner);
        assert!(error.backtrace().is_some());

        let error = ArcBacktraceFrom::from(Inner);
        assert!(error.backtrace().is_some());
    }
}

#[cfg(thiserror_nightly_testing)]
pub mod enums {
    use super::Inner;
    use std::backtrace::Backtrace;
    use std::error::Error;
    use std::sync::Arc;
    use thiserror::Error;

    #[derive(Error, Debug)]
    pub enum PlainBacktrace {
        #[error("...")]
        Test { backtrace: Backtrace },
    }

    #[derive(Error, Debug)]
    pub enum ExplicitBacktrace {
        #[error("...")]
        Test {
            #[backtrace]
            backtrace: Backtrace,
        },
    }

    #[derive(Error, Debug)]
    pub enum OptBacktrace {
        #[error("...")]
        Test {
            #[backtrace]
            backtrace: Option<Backtrace>,
        },
    }

    #[derive(Error, Debug)]
    pub enum ArcBacktrace {
        #[error("...")]
        Test {
            #[backtrace]
            backtrace: Arc<Backtrace>,
        },
    }

    #[derive(Error, Debug)]
    pub enum BacktraceFrom {
        #[error("...")]
        Test {
            #[from]
            source: Inner,
            #[backtrace]
            backtrace: Backtrace,
        },
    }

    #[derive(Error, Debug)]
    pub enum OptBacktraceFrom {
        #[error("...")]
        Test {
            #[from]
            source: Inner,
            #[backtrace]
            backtrace: Option<Backtrace>,
        },
    }

    #[derive(Error, Debug)]
    pub enum ArcBacktraceFrom {
        #[error("...")]
        Test {
            #[from]
            source: Inner,
            #[backtrace]
            backtrace: Arc<Backtrace>,
        },
    }

    #[test]
    fn test_backtrace() {
        let error = PlainBacktrace::Test {
            backtrace: Backtrace::capture(),
        };
        assert!(error.backtrace().is_some());

        let error = ExplicitBacktrace::Test {
            backtrace: Backtrace::capture(),
        };
        assert!(error.backtrace().is_some());

        let error = OptBacktrace::Test {
            backtrace: Some(Backtrace::capture()),
        };
        assert!(error.backtrace().is_some());

        let error = ArcBacktrace::Test {
            backtrace: Arc::new(Backtrace::capture()),
        };
        assert!(error.backtrace().is_some());

        let error = BacktraceFrom::from(Inner);
        assert!(error.backtrace().is_some());

        let error = OptBacktraceFrom::from(Inner);
        assert!(error.backtrace().is_some());

        let error = ArcBacktraceFrom::from(Inner);
        assert!(error.backtrace().is_some());
    }
}

#[test]
#[cfg_attr(not(thiserror_nightly_testing), ignore)]
fn test_backtrace() {}
<|endoftext|>


<commit_before>use std;
// use std::cell;
use std::cmp;
use std::ptr;
use std::sync::Arc;

use parking_lot::RwLock;

use io;

const KB_SIZE: usize = 1024;
pub type SharedMemoryNode = RwLock<[u8; KB_SIZE]>;

pub(crate) trait MemoryBlock {
    fn get_bytes(&self) -> u32;
    unsafe fn read_to_ptr(&self, offset: usize, buf: *mut u8, buf_size: usize);
    unsafe fn write_from_ptr(&mut self, offset: usize, buf: *const u8, buf_size: usize);
}

pub struct UniqueMemoryBlock(Vec<u8>);
impl UniqueMemoryBlock {
    pub fn new(kbs: usize) -> UniqueMemoryBlock {
        UniqueMemoryBlock(vec![0u8; kbs*KB_SIZE])
    }
}
impl MemoryBlock for UniqueMemoryBlock {
    fn get_bytes(&self) -> u32 {
        self.0.len() as u32
    }

    unsafe fn read_to_ptr(&self, offset: usize, buf: *mut u8, buf_size: usize) {
        let vec = &self.0;
        assert!(offset + buf_size <= vec.len());
        ptr::copy_nonoverlapping(vec.as_ptr().offset(offset as isize), buf, buf_size);
    }

    unsafe fn write_from_ptr(&mut self, offset: usize, buf: *const u8, buf_size: usize) {
        let vec = &mut self.0;
        assert!(offset + buf_size <= vec.len());
        ptr::copy_nonoverlapping(buf, vec.as_mut_ptr().offset(offset as isize), buf_size);
    }
}

#[derive(Clone)]
pub struct SharedMemoryBlock(Arc<Vec<SharedMemoryNode>>);
impl SharedMemoryBlock {
    pub fn new(kbs: usize) -> SharedMemoryBlock {
        let mut inner: Vec<SharedMemoryNode> = Vec::new();
        for _ in 0..kbs {
            inner.push(RwLock::new([0; KB_SIZE]))
        }

        SharedMemoryBlock(Arc::new(inner))
    }
}
impl MemoryBlock for SharedMemoryBlock {
    fn get_bytes(&self) -> u32 {
        let nodes = &self.0;
        (nodes.len() * KB_SIZE) as u32
    }

    unsafe fn read_to_ptr(&self, offset: usize, buf: *mut u8, buf_size: usize) {
        let nodes = &self.0;
        let mut buf_remaining = buf_size;
        let mut node_index = offset / KB_SIZE;
        let mut node_pos = offset % KB_SIZE;

        while buf_remaining > 0 {
            assert!(node_index < nodes.len());
            let copy_amount = cmp::min(KB_SIZE - node_pos, buf_remaining);
            {
                let node_ptr = nodes[node_index].read().as_ptr();
                let buf_pos = buf_size - buf_remaining;
                ptr::copy_nonoverlapping(node_ptr.offset(node_pos as isize), buf.offset(buf_pos as isize), copy_amount);
            }
            buf_remaining -= copy_amount;
            node_index += 1;
            node_pos = 0;
        }
    }

    unsafe fn write_from_ptr(&mut self, offset: usize, buf: *const u8, buf_size: usize) {
        let nodes = &self.0;
        let mut buf_remaining = buf_size;
        let mut node_index = offset / KB_SIZE;
        let mut node_pos = offset % KB_SIZE;

        while buf_remaining > 0 {
            assert!(node_index < nodes.len());
            let copy_amount = cmp::min(KB_SIZE - node_pos, buf_remaining);
            {
                let node_ptr = nodes[node_index].write().as_mut_ptr();
                let buf_pos = buf_size - buf_remaining;
                ptr::copy_nonoverlapping(buf.offset(buf_pos as isize), node_ptr.offset(node_pos as isize), copy_amount);
            }
            buf_remaining -= copy_amount;
            node_index += 1;
            node_pos = 0;
        }
    }
}

pub enum AddressBlock {
    UniqueRam(UniqueMemoryBlock),
    SharedRam(SharedMemoryBlock),
    Io9(io::IoRegsArm9),
    Io11(io::IoRegsArm11),
    IoPriv11(io::IoRegsArm11Priv),
    IoShared(io::IoRegsShared),
}
impl MemoryBlock for AddressBlock {
    fn get_bytes(&self) -> u32 {
        match *self {
            AddressBlock::UniqueRam(ref inner) => inner.get_bytes(),
            AddressBlock::SharedRam(ref inner) => inner.get_bytes(),
            AddressBlock::Io9(ref inner) => inner.get_bytes(),
            AddressBlock::Io11(ref inner) => inner.get_bytes(),
            AddressBlock::IoPriv11(ref inner) => inner.get_bytes(),
            AddressBlock::IoShared(ref inner) => inner.get_bytes(),
        }
    }

    unsafe fn read_to_ptr(&self, offset: usize, buf: *mut u8, buf_size: usize) {
        match *self {
            AddressBlock::UniqueRam(ref inner) => inner.read_to_ptr(offset, buf, buf_size),
            AddressBlock::SharedRam(ref inner) => inner.read_to_ptr(offset, buf, buf_size),
            AddressBlock::Io9(ref inner) => inner.read_to_ptr(offset, buf, buf_size),
            AddressBlock::Io11(ref inner) => inner.read_to_ptr(offset, buf, buf_size),
            AddressBlock::IoPriv11(ref inner) => inner.read_to_ptr(offset, buf, buf_size),
            AddressBlock::IoShared(ref inner) => inner.read_to_ptr(offset, buf, buf_size)
        }
    }

    unsafe fn write_from_ptr(&mut self, offset: usize, buf: *const u8, buf_size: usize) {
        match *self {
            AddressBlock::UniqueRam(ref mut inner) => inner.write_from_ptr(offset, buf, buf_size),
            AddressBlock::SharedRam(ref mut inner) => inner.write_from_ptr(offset, buf, buf_size),
            AddressBlock::Io9(ref mut inner) => inner.write_from_ptr(offset, buf, buf_size),
            AddressBlock::Io11(ref mut inner) => inner.write_from_ptr(offset, buf, buf_size),
            AddressBlock::IoPriv11(ref mut inner) => inner.write_from_ptr(offset, buf, buf_size),
            AddressBlock::IoShared(ref mut inner) => inner.write_from_ptr(offset, buf, buf_size)
        }
    }
}

pub struct AddressBlockHandle(u32);

pub struct MemController {
    regions: Vec<(u32, AddressBlock)>,
}

impl MemController {
    pub fn new() -> MemController {
        MemController {
            regions: Vec::new(),
        }
    }

    fn search_region(&self, address: u32) -> Result<usize, usize> {
        self.regions.binary_search_by(|&(addr, _)| addr.cmp(&address))
    }

    fn match_address<'a>(&'a self, address: u32) -> Option<usize> {
        let index = match self.search_region(address) {
            Ok(a) => a,
            Err(a) => a - 1,
        };
        if index > self.regions.len() {
            return None;
        }

        let (block_addr, ref block) = self.regions[index];
        if address - block_addr < block.get_bytes() {
            return Some(index);
        }
        None
    }

    pub fn map_region(&mut self, address: u32, region: AddressBlock) -> AddressBlockHandle {
        let insert_index = self.search_region(address).unwrap_err();
        self.regions.insert(insert_index, (address, region));
        AddressBlockHandle(address)
    }

    pub(crate) fn region(&self, handle: &AddressBlockHandle) -> &AddressBlock {
        let index = self.search_region(handle.0)
            .expect("Attempted to find region from non-existant handle!");
        &self.regions[index].1
    }

    pub(crate) fn _region_mut(&mut self, handle: &AddressBlockHandle) -> &mut AddressBlock {
        let index = self.search_region(handle.0)
            .expect("Attempted to find region from non-existant handle!");
        &mut self.regions[index].1
    }

    pub fn read<T: Copy>(&self, addr: u32) -> T {
        let block_index = self.match_address(addr)
            .unwrap_or_else(|| panic!("Could not match address 0x{:X}", addr));
        let (block_addr, ref block) = self.regions[block_index];

        unsafe {
            let mut t: T = std::mem::uninitialized();
            block.read_to_ptr((addr - block_addr) as usize, std::mem::transmute(&mut t), std::mem::size_of::<T>());
            t
        }
    }

    #[inline]
    fn try_read_buf(&self, addr: u32, buf: &mut [u8], debug: bool) -> Result<(), String> {
        let block_index = self.match_address(addr)
            .ok_or(format!("Could not match address 0x{:X}", addr))?;
        let (block_addr, ref block) = self.regions[block_index];

        match (debug, block) {
            | (true, AddressBlock::Io9(_))
            | (true, AddressBlock::IoShared(_)) => return Err(format!("Cannot issue debug read for IO address 0x{:X}", addr)),
            (_, block) => unsafe {
                block.read_to_ptr((addr - block_addr) as usize, buf.as_mut_ptr(), buf.len());
            }
        }
        Ok(())
    }

    pub fn read_buf(&self, addr: u32, buf: &mut [u8]) {
        self.try_read_buf(addr, buf, false).unwrap();
    }

    pub fn debug_read_buf(&self, addr: u32, buf: &mut [u8]) -> Result<(), String> {
        self.try_read_buf(addr, buf, true)
    }

    pub fn write<T: Copy>(&mut self, addr: u32, data: T) {
        let block_index = self.match_address(addr)
            .unwrap_or_else(|| panic!("Could not match address 0x{:X}", addr));
        let (block_addr, ref mut block) = self.regions[block_index];

        unsafe {
            block.write_from_ptr((addr - block_addr) as usize, std::mem::transmute(&data), std::mem::size_of::<T>());
        }
    }

    pub fn write_buf(&mut self, addr: u32, buf: &[u8]) {
        let block_index = self.match_address(addr)
            .unwrap_or_else(|| panic!("Could not match address 0x{:X}", addr));
        let (block_addr, ref mut block) = self.regions[block_index];

        unsafe {
            block.write_from_ptr((addr - block_addr) as usize, buf.as_ptr(), buf.len());
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::io::Write;

    #[test]
    fn write_intra_block() {
        let block = SharedMemoryBlock::new(1);
        assert_eq!(block.get_bytes(), 0x400);
        let nodes = &block.0;

        // Write data
        let bytes = [0xFFu8, 0x53u8, 0x28u8, 0xC6u8];
        unsafe {
            block.write_from_ptr(0x2C8, bytes.as_ptr(), bytes.len());
        }

        // Compare memory with data
        let block_mem = nodes[0].read();
        assert_eq!(block_mem[0x2C8..0x2CC], bytes[..]);
    }

    #[test]
    fn read_intra_block() {
        let block = SharedMemoryBlock::new(1);
        assert_eq!(block.get_bytes(), 0x400);
        let nodes = &block.0;

        // Write data directly to memory
        {
            let mut block_mem = nodes[0].write();
            (&mut block_mem[0x2C8..0x2CC]).write_all(&[0xFFu8, 0x53u8, 0x28u8, 0xC6u8]);
        }

        // Read memory to buffer
        let mut buf = [0u8; 4];
        unsafe {
            block.read_to_ptr(0x2C8, buf.as_mut_ptr(), buf.len());
        }

        // Compare memory and buffer
        let block_mem = nodes[0].read();
        assert_eq!(block_mem[0x2C8..0x2CC], buf[..]);
    }

    #[test]
    fn write_inter_block() {
        let block = SharedMemoryBlock::new(2);
        assert_eq!(block.get_bytes(), 0x800);
        let nodes = &block.0;

        // Write data
        let bytes = [0xFFu8, 0x53u8, 0x28u8, 0xC6u8];
        unsafe {
            block.write_from_ptr(0x3FE, bytes.as_ptr(), bytes.len());
        }

        // Compare memory with data
        let block_mem0 = nodes[0].read();
        let block_mem1 = nodes[1].read();
        assert_eq!(block_mem0[0x3FE..0x400], bytes[0..2]);
        assert_eq!(block_mem1[0x0..0x2], bytes[2..4]);
    }

    #[test]
    fn read_inter_block() {
        let block = SharedMemoryBlock::new(2);
        assert_eq!(block.get_bytes(), 0x800);
        let nodes = &block.0;

        // Write data directly to memory
        {
            let mut block_mem0 = nodes[0].write();
            let mut block_mem1 = nodes[1].write();
            (&mut block_mem0[0x3FE..0x400]).write_all(&[0xFFu8, 0x53u8]);
            (&mut block_mem1[0x0..0x2]).write_all(&[0x28u8, 0xC6u8]);
        }

        // Read memory to buffer
        let mut buf = [0u8; 4];
        unsafe {
            block.read_to_ptr(0x3FE, buf.as_mut_ptr(), buf.len());
        }

        // Compare memory and buffer
        let block_mem0 = nodes[0].read();
        let block_mem1 = nodes[1].read();
        assert_eq!(block_mem0[0x3FE..0x400], buf[0..2]);
        assert_eq!(block_mem1[0x0..0x2], buf[2..4]);
    }
}
<commit_msg>mem: Use BTreeMap instead of binary searched vec for regions<commit_after>use std;
use std::cmp;
use std::ptr;
use std::sync::Arc;

use std::collections::BTreeMap;

use parking_lot::RwLock;

use io;

const KB_SIZE: usize = 1024;
pub type SharedMemoryNode = RwLock<[u8; KB_SIZE]>;

pub(crate) trait MemoryBlock {
    fn get_bytes(&self) -> u32;
    unsafe fn read_to_ptr(&self, offset: usize, buf: *mut u8, buf_size: usize);
    unsafe fn write_from_ptr(&mut self, offset: usize, buf: *const u8, buf_size: usize);
}

pub struct UniqueMemoryBlock(Vec<u8>);
impl UniqueMemoryBlock {
    pub fn new(kbs: usize) -> UniqueMemoryBlock {
        UniqueMemoryBlock(vec![0u8; kbs*KB_SIZE])
    }
}
impl MemoryBlock for UniqueMemoryBlock {
    fn get_bytes(&self) -> u32 {
        self.0.len() as u32
    }

    unsafe fn read_to_ptr(&self, offset: usize, buf: *mut u8, buf_size: usize) {
        let vec = &self.0;
        assert!(offset + buf_size <= vec.len());
        ptr::copy_nonoverlapping(vec.as_ptr().offset(offset as isize), buf, buf_size);
    }

    unsafe fn write_from_ptr(&mut self, offset: usize, buf: *const u8, buf_size: usize) {
        let vec = &mut self.0;
        assert!(offset + buf_size <= vec.len());
        ptr::copy_nonoverlapping(buf, vec.as_mut_ptr().offset(offset as isize), buf_size);
    }
}

#[derive(Clone)]
pub struct SharedMemoryBlock(Arc<Vec<SharedMemoryNode>>);
impl SharedMemoryBlock {
    pub fn new(kbs: usize) -> SharedMemoryBlock {
        let mut inner: Vec<SharedMemoryNode> = Vec::new();
        for _ in 0..kbs {
            inner.push(RwLock::new([0; KB_SIZE]))
        }

        SharedMemoryBlock(Arc::new(inner))
    }
}
impl MemoryBlock for SharedMemoryBlock {
    fn get_bytes(&self) -> u32 {
        let nodes = &self.0;
        (nodes.len() * KB_SIZE) as u32
    }

    unsafe fn read_to_ptr(&self, offset: usize, buf: *mut u8, buf_size: usize) {
        let nodes = &self.0;
        let mut buf_remaining = buf_size;
        let mut node_index = offset / KB_SIZE;
        let mut node_pos = offset % KB_SIZE;

        while buf_remaining > 0 {
            assert!(node_index < nodes.len());
            let copy_amount = cmp::min(KB_SIZE - node_pos, buf_remaining);
            {
                let node_ptr = nodes[node_index].read().as_ptr();
                let buf_pos = buf_size - buf_remaining;
                ptr::copy_nonoverlapping(node_ptr.offset(node_pos as isize), buf.offset(buf_pos as isize), copy_amount);
            }
            buf_remaining -= copy_amount;
            node_index += 1;
            node_pos = 0;
        }
    }

    unsafe fn write_from_ptr(&mut self, offset: usize, buf: *const u8, buf_size: usize) {
        let nodes = &self.0;
        let mut buf_remaining = buf_size;
        let mut node_index = offset / KB_SIZE;
        let mut node_pos = offset % KB_SIZE;

        while buf_remaining > 0 {
            assert!(node_index < nodes.len());
            let copy_amount = cmp::min(KB_SIZE - node_pos, buf_remaining);
            {
                let node_ptr = nodes[node_index].write().as_mut_ptr();
                let buf_pos = buf_size - buf_remaining;
                ptr::copy_nonoverlapping(buf.offset(buf_pos as isize), node_ptr.offset(node_pos as isize), copy_amount);
            }
            buf_remaining -= copy_amount;
            node_index += 1;
            node_pos = 0;
        }
    }
}

pub enum AddressBlock {
    UniqueRam(UniqueMemoryBlock),
    SharedRam(SharedMemoryBlock),
    Io9(io::IoRegsArm9),
    Io11(io::IoRegsArm11),
    IoPriv11(io::IoRegsArm11Priv),
    IoShared(io::IoRegsShared),
}
impl MemoryBlock for AddressBlock {
    fn get_bytes(&self) -> u32 {
        match *self {
            AddressBlock::UniqueRam(ref inner) => inner.get_bytes(),
            AddressBlock::SharedRam(ref inner) => inner.get_bytes(),
            AddressBlock::Io9(ref inner) => inner.get_bytes(),
            AddressBlock::Io11(ref inner) => inner.get_bytes(),
            AddressBlock::IoPriv11(ref inner) => inner.get_bytes(),
            AddressBlock::IoShared(ref inner) => inner.get_bytes(),
        }
    }

    unsafe fn read_to_ptr(&self, offset: usize, buf: *mut u8, buf_size: usize) {
        match *self {
            AddressBlock::UniqueRam(ref inner) => inner.read_to_ptr(offset, buf, buf_size),
            AddressBlock::SharedRam(ref inner) => inner.read_to_ptr(offset, buf, buf_size),
            AddressBlock::Io9(ref inner) => inner.read_to_ptr(offset, buf, buf_size),
            AddressBlock::Io11(ref inner) => inner.read_to_ptr(offset, buf, buf_size),
            AddressBlock::IoPriv11(ref inner) => inner.read_to_ptr(offset, buf, buf_size),
            AddressBlock::IoShared(ref inner) => inner.read_to_ptr(offset, buf, buf_size)
        }
    }

    unsafe fn write_from_ptr(&mut self, offset: usize, buf: *const u8, buf_size: usize) {
        match *self {
            AddressBlock::UniqueRam(ref mut inner) => inner.write_from_ptr(offset, buf, buf_size),
            AddressBlock::SharedRam(ref mut inner) => inner.write_from_ptr(offset, buf, buf_size),
            AddressBlock::Io9(ref mut inner) => inner.write_from_ptr(offset, buf, buf_size),
            AddressBlock::Io11(ref mut inner) => inner.write_from_ptr(offset, buf, buf_size),
            AddressBlock::IoPriv11(ref mut inner) => inner.write_from_ptr(offset, buf, buf_size),
            AddressBlock::IoShared(ref mut inner) => inner.write_from_ptr(offset, buf, buf_size)
        }
    }
}

pub struct AddressBlockHandle(u32);

pub struct MemController {
    regions: BTreeMap<u32, AddressBlock>,
}

impl MemController {
    pub fn new() -> MemController {
        MemController {
            regions: BTreeMap::new(),
        }
    }

    fn match_address(&self, address: u32) -> Option<(u32, &AddressBlock)> {
        let matched = self.regions.range(..=address).next_back()?;
        let (block_addr, block) = matched;
        if address - block_addr < block.get_bytes() {
            return Some((*block_addr, block));
        }
        None
    }

    fn match_address_mut(&mut self, address: u32) -> Option<(u32, &mut AddressBlock)> {
        let matched = self.regions.range_mut(..=address).next_back()?;
        let (block_addr, block) = matched;
        if address - block_addr < block.get_bytes() {
            return Some((*block_addr, block));
        }
        None
    }

    pub fn map_region(&mut self, address: u32, region: AddressBlock) -> AddressBlockHandle {
        self.regions.insert(address, region);
        AddressBlockHandle(address)
    }

    pub(crate) fn region(&self, handle: &AddressBlockHandle) -> &AddressBlock {
        self.regions.get(&handle.0)
            .expect("Attempted to find region from non-existant handle!")
    }

    pub(crate) fn _region_mut(&mut self, handle: &AddressBlockHandle) -> &mut AddressBlock {
        self.regions.get_mut(&handle.0)
            .expect("Attempted to find region from non-existant handle!")
    }

    pub fn read<T: Copy>(&self, addr: u32) -> T {
        let (block_addr, ref block) = self.match_address(addr)
            .unwrap_or_else(|| panic!("Could not match address 0x{:X}", addr));

        unsafe {
            let mut t: T = std::mem::uninitialized();
            block.read_to_ptr((addr - block_addr) as usize, std::mem::transmute(&mut t), std::mem::size_of::<T>());
            t
        }
    }

    #[inline]
    fn try_read_buf(&self, addr: u32, buf: &mut [u8], debug: bool) -> Result<(), String> {
        let (block_addr, ref block) = self.match_address(addr)
            .ok_or(format!("Could not match address 0x{:X}", addr))?;

        match (debug, block) {
            | (true, AddressBlock::Io9(_))
            | (true, AddressBlock::IoShared(_)) => return Err(format!("Cannot issue debug read for IO address 0x{:X}", addr)),
            (_, block) => unsafe {
                block.read_to_ptr((addr - block_addr) as usize, buf.as_mut_ptr(), buf.len());
            }
        }
        Ok(())
    }

    pub fn read_buf(&self, addr: u32, buf: &mut [u8]) {
        self.try_read_buf(addr, buf, false).unwrap();
    }

    pub fn debug_read_buf(&self, addr: u32, buf: &mut [u8]) -> Result<(), String> {
        self.try_read_buf(addr, buf, true)
    }

    pub fn write<T: Copy>(&mut self, addr: u32, data: T) {
        let (block_addr, ref mut block) = self.match_address_mut(addr)
            .unwrap_or_else(|| panic!("Could not match address 0x{:X}", addr));

        unsafe {
            block.write_from_ptr((addr - block_addr) as usize, std::mem::transmute(&data), std::mem::size_of::<T>());
        }
    }

    pub fn write_buf(&mut self, addr: u32, buf: &[u8]) {
        let (block_addr, ref mut block) = self.match_address_mut(addr)
            .unwrap_or_else(|| panic!("Could not match address 0x{:X}", addr));

        unsafe {
            block.write_from_ptr((addr - block_addr) as usize, buf.as_ptr(), buf.len());
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::io::Write;

    #[test]
    fn write_intra_block() {
        let block = SharedMemoryBlock::new(1);
        assert_eq!(block.get_bytes(), 0x400);
        let nodes = &block.0;

        // Write data
        let bytes = [0xFFu8, 0x53u8, 0x28u8, 0xC6u8];
        unsafe {
            block.write_from_ptr(0x2C8, bytes.as_ptr(), bytes.len());
        }

        // Compare memory with data
        let block_mem = nodes[0].read();
        assert_eq!(block_mem[0x2C8..0x2CC], bytes[..]);
    }

    #[test]
    fn read_intra_block() {
        let block = SharedMemoryBlock::new(1);
        assert_eq!(block.get_bytes(), 0x400);
        let nodes = &block.0;

        // Write data directly to memory
        {
            let mut block_mem = nodes[0].write();
            (&mut block_mem[0x2C8..0x2CC]).write_all(&[0xFFu8, 0x53u8, 0x28u8, 0xC6u8]);
        }

        // Read memory to buffer
        let mut buf = [0u8; 4];
        unsafe {
            block.read_to_ptr(0x2C8, buf.as_mut_ptr(), buf.len());
        }

        // Compare memory and buffer
        let block_mem = nodes[0].read();
        assert_eq!(block_mem[0x2C8..0x2CC], buf[..]);
    }

    #[test]
    fn write_inter_block() {
        let block = SharedMemoryBlock::new(2);
        assert_eq!(block.get_bytes(), 0x800);
        let nodes = &block.0;

        // Write data
        let bytes = [0xFFu8, 0x53u8, 0x28u8, 0xC6u8];
        unsafe {
            block.write_from_ptr(0x3FE, bytes.as_ptr(), bytes.len());
        }

        // Compare memory with data
        let block_mem0 = nodes[0].read();
        let block_mem1 = nodes[1].read();
        assert_eq!(block_mem0[0x3FE..0x400], bytes[0..2]);
        assert_eq!(block_mem1[0x0..0x2], bytes[2..4]);
    }

    #[test]
    fn read_inter_block() {
        let block = SharedMemoryBlock::new(2);
        assert_eq!(block.get_bytes(), 0x800);
        let nodes = &block.0;

        // Write data directly to memory
        {
            let mut block_mem0 = nodes[0].write();
            let mut block_mem1 = nodes[1].write();
            (&mut block_mem0[0x3FE..0x400]).write_all(&[0xFFu8, 0x53u8]);
            (&mut block_mem1[0x0..0x2]).write_all(&[0x28u8, 0xC6u8]);
        }

        // Read memory to buffer
        let mut buf = [0u8; 4];
        unsafe {
            block.read_to_ptr(0x3FE, buf.as_mut_ptr(), buf.len());
        }

        // Compare memory and buffer
        let block_mem0 = nodes[0].read();
        let block_mem1 = nodes[1].read();
        assert_eq!(block_mem0[0x3FE..0x400], buf[0..2]);
        assert_eq!(block_mem1[0x0..0x2], buf[2..4]);
    }
}
<|endoftext|>


<commit_before>// extern crate bindgen;
#[cfg(target_env = "msvc")]
extern crate cc;
extern crate pkg_config;
extern crate num_cpus;

use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

// Automatically write bindings to libsass
//#[allow(dead_code)]
//fn write_bindings() {
//    let bindings = bindgen::Builder::default()
//        .header("libsass/include/sass.h")
//        .clang_arg("-Ilibsass/include")
//        // To avoid a test failing
//        .blacklist_type("max_align_t")
//        // we do static linking so it should be fine
//        // https://github.com/rust-lang/rust/issues/36927
//        .rustified_enum(".*")
//        .generate()
//        .expect("Unable to generate bindings");
//
//    // Write the bindings to the $OUT_DIR/bindings.rs file.
//    let out_path = PathBuf::from("src");
//    bindings
//        .write_to_file(out_path.join("bindings.rs"))
//        .expect("Couldn't write bindings!");
//}

macro_rules! t {
    ($e:expr) => (match $e {
        Ok(n) => n,
        Err(e) => panic!("\n{} failed with {}\n", stringify!($e), e),
    })
}

fn cp_r(dir: &Path, dest: &Path) {
    for entry in t!(fs::read_dir(dir)) {
        let entry = t!(entry);
        let path = entry.path();
        let dst = dest.join(path.file_name().expect("Failed to get filename of path"));
        if t!(fs::metadata(&path)).is_file() {
            t!(fs::copy(path, dst));
        } else {
            t!(fs::create_dir_all(&dst));
            cp_r(&path, &dst);
        }
    }
}

fn get_libsass_folder() -> PathBuf {
    env::current_dir().expect("Failed to get the current directory").join("libsass")
}

#[cfg(target_os = "linux")]
fn _compile(libprobe: fn(&str) -> bool) {
    if libprobe("stdc++") {
        println!("cargo:rustc-link-lib=dylib=stdc++");
    } else if libprobe("c++_shared") {
        println!("cargo:rustc-link-lib=dylib=c++_shared");
    } else if libprobe("c++") {
        println!("cargo:rustc-link-lib=dylib=c++");
    } else {
        panic!("no c++ library found");
    }
}

#[cfg(not(target_os = "linux"))]
fn _compile(libprobe: fn(&str) -> bool) {
    if libprobe("c++_shared") {
        println!("cargo:rustc-link-lib=dylib=c++_shared");
    } else if libprobe("c++") {
        println!("cargo:rustc-link-lib=dylib=c++");
    } else if libprobe("stdc++") {
        println!("cargo:rustc-link-lib=dylib=stdc++");
    } else {
        panic!("no c++ library found");
    }
}

// linux/unix
#[cfg(not(target_env = "msvc"))]
fn compile() {
    let target = env::var("TARGET").expect("TARGET not found");
    let src = get_libsass_folder();
    let dest = PathBuf::from(env::var_os("OUT_DIR").expect("OUT_DIR not found"));
    let build = dest.join("build");
    t!(fs::create_dir_all(&build));
    cp_r(&src, &build);
    let is_bsd = target.contains("dragonfly")
        || target.contains("freebsd")
        || target.contains("netbsd")
        || target.contains("openbsd");
    let jobs = env::var("MAKE_LIBSASS_JOBS").unwrap_or_else(|_| num_cpus::get().to_string());
    let r = Command::new(if is_bsd { "gmake" } else { "make" })
        .current_dir(&build)
        .args(&["--jobs", &jobs])
        .output()
        .expect("Failed to run make/gmake. Do you have it installed?");

    if !r.status.success() {
        let err = String::from_utf8_lossy(&r.stderr);
        let out = String::from_utf8_lossy(&r.stdout);
        panic!("Build error:\nSTDERR:{}\nSTDOUT:{}", err, out);
    }

    println!(
        "cargo:rustc-link-search=native={}",
        build.join("lib").display()
    );
    println!("cargo:rustc-link-lib=static=sass");

    let libprobe = | lib: &str | -> bool {
      Command::new("cc")
        .arg("-xc++")
        .arg("-o/dev/null")
        .arg(format!("-l{}",lib))
        .arg("-shared")
        .stderr(Stdio::null())
        .status()
        .expect("Failed to run cc. Do you have it installed?")
        .success()
    };
    _compile(libprobe);
}

// windows
#[cfg(target_env = "msvc")]
fn compile() {
    let src = get_libsass_folder();
    let target = env::var("TARGET").expect("TARGET not found in environment");
    let msvc_platform = if target.contains("x86_64") {
        "Win64"
    } else {
        "Win32"
    };
    let dest = PathBuf::from(env::var_os("OUT_DIR").expect("OUT_DIR not found in environment"));
    let build = dest.join("build");
    t!(fs::create_dir_all(&build));
    cp_r(&src, &build);

    // Find an instance of devenv.exe from Visual Studio IDE in order to upgrade
    // libsass.sln to the current available IDE. Do nothing if no devenv.exe found
    let d = cc::windows_registry::find(target.as_str(), "devenv.exe");
    if let Some(mut d) = d {
        let d = d
            .args(&["/upgrade", "win\\libsass.sln"])
            .current_dir(&build)
            .output()
            .expect("Failed to run devenv. Do you have it installed?");
        if !d.status.success() {
            let err = String::from_utf8_lossy(&d.stderr);
            let out = String::from_utf8_lossy(&d.stdout);
            println!("Upgrade error:\nSTDERR:{}\nSTDOUT:{}", err, out);
        }
    }

    let search = Command::new("where")
        .args(&["msbuild.exe"])
        .output()
        .expect("Failed to run where: Could not search for msbuild.exe on path.");
    let mut msbuild = if search.status.success() {
        Command::new("msbuild.exe")
    } else {
        cc::windows_registry::find(target.as_str(), "msbuild.exe")
            .expect("Could not find msbuild.exe on the registry")
    };

    let jobs = env::var("MAKE_LIBSASS_JOBS").unwrap_or_else(|_| num_cpus::get().to_string());

    let r = msbuild
        .args(&[
            "win\\libsass.sln",
            "/p:LIBSASS_STATIC_LIB=1",
            "/p:Configuration=Release",
            "/p:WholeProgramOptimization=false",
            format!("/m:{}", jobs).as_str(),
            format!("/p:Platform={}", msvc_platform).as_str(),
        ])
        .current_dir(&build)
        .output()
        .expect("Failed to run msbuild. Do you have it installed?");

    if !r.status.success() {
        let err = String::from_utf8_lossy(&r.stderr);
        let out = String::from_utf8_lossy(&r.stdout);
        panic!("Build error:\nSTDERR:{}\nSTDOUT:{}", err, out);
    }

    println!(
        "cargo:rustc-link-search=native={}",
        build.join("win").join("bin").display()
    );
    println!("cargo:rustc-link-lib=static=libsass");
}

fn main() {
    // Uncomment the line below to generate bindings. Doesn't work on CI as it
    // requires additional tooling
    // write_bindings();

    // Is it already built?
    if let Ok(_) = pkg_config::find_library("libsass") {
        return;
    }

    compile();
}
<commit_msg>Add error message if libsass source is not available<commit_after>// extern crate bindgen;
#[cfg(target_env = "msvc")]
extern crate cc;
extern crate pkg_config;
extern crate num_cpus;

use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

// Automatically write bindings to libsass
//#[allow(dead_code)]
//fn write_bindings() {
//    let bindings = bindgen::Builder::default()
//        .header("libsass/include/sass.h")
//        .clang_arg("-Ilibsass/include")
//        // To avoid a test failing
//        .blacklist_type("max_align_t")
//        // we do static linking so it should be fine
//        // https://github.com/rust-lang/rust/issues/36927
//        .rustified_enum(".*")
//        .generate()
//        .expect("Unable to generate bindings");
//
//    // Write the bindings to the $OUT_DIR/bindings.rs file.
//    let out_path = PathBuf::from("src");
//    bindings
//        .write_to_file(out_path.join("bindings.rs"))
//        .expect("Couldn't write bindings!");
//}

macro_rules! t {
    ($e:expr) => (match $e {
        Ok(n) => n,
        Err(e) => panic!("\n{} failed with {}\n", stringify!($e), e),
    })
}

fn cp_r(dir: &Path, dest: &Path) {
    for entry in t!(fs::read_dir(dir)) {
        let entry = t!(entry);
        let path = entry.path();
        let dst = dest.join(path.file_name().expect("Failed to get filename of path"));
        if t!(fs::metadata(&path)).is_file() {
            t!(fs::copy(path, dst));
        } else {
            t!(fs::create_dir_all(&dst));
            cp_r(&path, &dst);
        }
    }
}

fn get_libsass_folder() -> PathBuf {
    env::current_dir().expect("Failed to get the current directory").join("libsass")
}

#[cfg(target_os = "linux")]
fn _compile(libprobe: fn(&str) -> bool) {
    if libprobe("stdc++") {
        println!("cargo:rustc-link-lib=dylib=stdc++");
    } else if libprobe("c++_shared") {
        println!("cargo:rustc-link-lib=dylib=c++_shared");
    } else if libprobe("c++") {
        println!("cargo:rustc-link-lib=dylib=c++");
    } else {
        panic!("no c++ library found");
    }
}

#[cfg(not(target_os = "linux"))]
fn _compile(libprobe: fn(&str) -> bool) {
    if libprobe("c++_shared") {
        println!("cargo:rustc-link-lib=dylib=c++_shared");
    } else if libprobe("c++") {
        println!("cargo:rustc-link-lib=dylib=c++");
    } else if libprobe("stdc++") {
        println!("cargo:rustc-link-lib=dylib=stdc++");
    } else {
        panic!("no c++ library found");
    }
}

// linux/unix
#[cfg(not(target_env = "msvc"))]
fn compile() {
    let target = env::var("TARGET").expect("TARGET not found");
    let src = get_libsass_folder();
    let dest = PathBuf::from(env::var_os("OUT_DIR").expect("OUT_DIR not found"));
    let build = dest.join("build");
    t!(fs::create_dir_all(&build));
    cp_r(&src, &build);
    let is_bsd = target.contains("dragonfly")
        || target.contains("freebsd")
        || target.contains("netbsd")
        || target.contains("openbsd");
    let jobs = env::var("MAKE_LIBSASS_JOBS").unwrap_or_else(|_| num_cpus::get().to_string());
    let r = Command::new(if is_bsd { "gmake" } else { "make" })
        .current_dir(&build)
        .args(&["--jobs", &jobs])
        .output()
        .expect("Failed to run make/gmake. Do you have it installed?");

    if !r.status.success() {
        let err = String::from_utf8_lossy(&r.stderr);
        let out = String::from_utf8_lossy(&r.stdout);
        panic!("Build error:\nSTDERR:{}\nSTDOUT:{}", err, out);
    }

    println!(
        "cargo:rustc-link-search=native={}",
        build.join("lib").display()
    );
    println!("cargo:rustc-link-lib=static=sass");

    let libprobe = | lib: &str | -> bool {
      Command::new("cc")
        .arg("-xc++")
        .arg("-o/dev/null")
        .arg(format!("-l{}",lib))
        .arg("-shared")
        .stderr(Stdio::null())
        .status()
        .expect("Failed to run cc. Do you have it installed?")
        .success()
    };
    _compile(libprobe);
}

// windows
#[cfg(target_env = "msvc")]
fn compile() {
    let src = get_libsass_folder();
    let target = env::var("TARGET").expect("TARGET not found in environment");
    let msvc_platform = if target.contains("x86_64") {
        "Win64"
    } else {
        "Win32"
    };
    let dest = PathBuf::from(env::var_os("OUT_DIR").expect("OUT_DIR not found in environment"));
    let build = dest.join("build");
    t!(fs::create_dir_all(&build));
    cp_r(&src, &build);

    // Find an instance of devenv.exe from Visual Studio IDE in order to upgrade
    // libsass.sln to the current available IDE. Do nothing if no devenv.exe found
    let d = cc::windows_registry::find(target.as_str(), "devenv.exe");
    if let Some(mut d) = d {
        let d = d
            .args(&["/upgrade", "win\\libsass.sln"])
            .current_dir(&build)
            .output()
            .expect("Failed to run devenv. Do you have it installed?");
        if !d.status.success() {
            let err = String::from_utf8_lossy(&d.stderr);
            let out = String::from_utf8_lossy(&d.stdout);
            println!("Upgrade error:\nSTDERR:{}\nSTDOUT:{}", err, out);
        }
    }

    let search = Command::new("where")
        .args(&["msbuild.exe"])
        .output()
        .expect("Failed to run where: Could not search for msbuild.exe on path.");
    let mut msbuild = if search.status.success() {
        Command::new("msbuild.exe")
    } else {
        cc::windows_registry::find(target.as_str(), "msbuild.exe")
            .expect("Could not find msbuild.exe on the registry")
    };

    let jobs = env::var("MAKE_LIBSASS_JOBS").unwrap_or_else(|_| num_cpus::get().to_string());

    let r = msbuild
        .args(&[
            "win\\libsass.sln",
            "/p:LIBSASS_STATIC_LIB=1",
            "/p:Configuration=Release",
            "/p:WholeProgramOptimization=false",
            format!("/m:{}", jobs).as_str(),
            format!("/p:Platform={}", msvc_platform).as_str(),
        ])
        .current_dir(&build)
        .output()
        .expect("Failed to run msbuild. Do you have it installed?");

    if !r.status.success() {
        let err = String::from_utf8_lossy(&r.stderr);
        let out = String::from_utf8_lossy(&r.stdout);
        panic!("Build error:\nSTDERR:{}\nSTDOUT:{}", err, out);
    }

    println!(
        "cargo:rustc-link-search=native={}",
        build.join("win").join("bin").display()
    );
    println!("cargo:rustc-link-lib=static=libsass");
}

fn main() {
    // Uncomment the line below to generate bindings. Doesn't work on CI as it
    // requires additional tooling
    // write_bindings();

    // Is it already built?
    if pkg_config::find_library("libsass").is_ok() {
        return;
    }

    if !get_libsass_folder().join("Makefile").exists() {
        eprintln!("
            Could not find libsass source
            Try running `git submodule update --init` if you are building from the repository
        ");
        std::process::exit(1);
    }

    compile();
}
<|endoftext|>


<commit_before>use {
    std::{
        io::SeekFrom,
        path::PathBuf,
        os::unix::{fs::PermissionsExt, io::FromRawFd},
    },
    filetime::FileTime,
    nix::libc::{O_CREAT, O_RDONLY, O_RDWR, O_TRUNC, O_WRONLY},
    async_trait::async_trait,
    tokio::{
        prelude::*,
        fs,
        stream::StreamExt,
        sync::{Mutex, RwLock},
    },
    rs9p::{*, srv::{srv_async, Fid, Filesystem}},
};

mod utils;
use crate::utils::*;

// Some clients will incorrectly set bits in 9p flags that don't make sense.
// For exmaple, the linux 9p kernel client propagates O_DIRECT to TCREATE and TOPEN
// and from there to the server.
// Processes on client machines set O_DIRECT to bypass the cache, but if
// the server uses O_DIRECT in the open or create, then subsequent server
// write and read system calls will fail, as O_DIRECT requires at minimum 512
// byte aligned data, and the data is almost always not aligned.
// While the linux kernel client is arguably broken, we won't be able
// to fix every kernel out there, and this is surely not the only buggy client
// we will see.
// The fix is to enumerate the set of flags we support and then and that with
// the flags received in a TCREATE or TOPEN. This nicely fixes a real problem
// we are seeing with a file system benchmark.
const UNIX_FLAGS: u32 = (O_WRONLY | O_RDONLY | O_RDWR | O_CREAT | O_TRUNC) as u32;

#[derive(Default)]
struct UnpfsFid {
    realpath: RwLock<PathBuf>,
    file: Mutex<Option<fs::File>>,
}

#[derive(Clone)]
struct Unpfs {
    realroot: PathBuf,
}

#[async_trait]
impl Filesystem for Unpfs {
    type Fid = UnpfsFid;

    async fn rattach(
        &self,
        fid: &Fid<Self::Fid>,
        _afid: Option<&Fid<Self::Fid>>,
        _uname: &str,
        _aname: &str,
        _n_uname: u32,
    ) -> Result<Fcall> {
        {
            let mut realpath = fid.aux.realpath.write().await;
            *realpath = PathBuf::from(&self.realroot);
        }

        Ok(Fcall::Rattach {
            qid: get_qid(&self.realroot).await?,
        })
    }

    async fn rwalk(
        &self,
        fid: &Fid<Self::Fid>,
        newfid: &Fid<Self::Fid>,
        wnames: &[String],
    ) -> Result<Fcall> {
        let mut wqids = Vec::new();
        let mut path = {
            let realpath = fid.aux.realpath.read().await;
            realpath.clone()
        };

        for (i, name) in wnames.iter().enumerate() {
            path.push(name);

            let qid = match get_qid(&path).await {
                Ok(qid) => qid,
                Err(e) => {
                    if i == 0 {
                        return Err(e);
                    } else {
                        break;
                    }
                }
            };

            wqids.push(qid);
        }

        {
            let mut new_realpath = newfid.aux.realpath.write().await;
            *new_realpath = path;
        }

        Ok(Fcall::Rwalk { wqids: wqids })
    }

    async fn rgetattr(&self, fid: &Fid<Self::Fid>, req_mask: GetattrMask) -> Result<Fcall> {
        let attr = {
            let realpath = fid.aux.realpath.read().await;
            fs::symlink_metadata(&*realpath).await?
        };

        Ok(Fcall::Rgetattr {
            valid: req_mask,
            qid: qid_from_attr(&attr),
            stat: From::from(attr),
        })
    }

    async fn rsetattr(
        &self,
        fid: &Fid<Self::Fid>,
        valid: SetattrMask,
        stat: &SetAttr,
    ) -> Result<Fcall> {
        let filepath = {
            let realpath = fid.aux.realpath.read().await;
            realpath.clone()
        };

        if valid.contains(SetattrMask::MODE) {
            fs::set_permissions(&filepath, PermissionsExt::from_mode(stat.mode)).await?;
        }

        if valid.intersects(SetattrMask::UID | SetattrMask::GID) {
            let uid = if valid.contains(SetattrMask::UID) {
                Some(nix::unistd::Uid::from_raw(stat.uid))
            } else {
                None
            };
            let gid = if valid.contains(SetattrMask::GID) {
                Some(nix::unistd::Gid::from_raw(stat.gid))
            } else {
                None
            };
            nix::unistd::chown(&filepath, uid, gid)?;
        }

        if valid.contains(SetattrMask::SIZE) {
            let _ = fs::File::open(&filepath).await?.set_len(stat.size).await?;
        }

        if valid.intersects(SetattrMask::ATIME_SET | SetattrMask::MTIME_SET) {
            let attr = fs::metadata(&filepath).await?;
            let atime = if valid.contains(SetattrMask::ATIME_SET) {
                FileTime::from_unix_time(stat.atime.sec as i64, stat.atime.nsec as u32)
            } else {
                FileTime::from_last_access_time(&attr)
            };

            let mtime = if valid.contains(SetattrMask::MTIME_SET) {
                FileTime::from_unix_time(stat.mtime.sec as i64, stat.mtime.nsec as u32)
            } else {
                FileTime::from_last_modification_time(&attr)
            };

            let _ = tokio::task::spawn_blocking(move || filetime::set_file_times(filepath, atime, mtime)).await;
        }

        Ok(Fcall::Rsetattr)
    }

    async fn rreadlink(&self, fid: &Fid<Self::Fid>) -> Result<Fcall> {
        let link = {
            let realpath = fid.aux.realpath.read().await;
            fs::read_link(&*realpath).await?
        };

        Ok(Fcall::Rreadlink {
            target: link.to_string_lossy().into_owned(),
        })
    }

    async fn rreaddir(&self, fid: &Fid<Self::Fid>, off: u64, count: u32) -> Result<Fcall> {
        let mut dirents = DirEntryData::new();

        let offset = if off == 0 {
            dirents.push(get_dirent_from(".", 0).await?);
            dirents.push(get_dirent_from("..", 1).await?);
            off
        } else {
            off - 1
        } as usize;

        let mut entries = {
            let realpath = fid.aux.realpath.read().await;
            fs::read_dir(&*realpath).await?.skip(offset)
        };

        let mut i = offset;
        while let Some(entry) = entries.next().await {
            let dirent = get_dirent(&entry?, 2 + i as u64).await?;
            if dirents.size() + dirent.size() > count {
                break;
            }
            dirents.push(dirent);
            i += 1;
        }

        Ok(Fcall::Rreaddir { data: dirents })
    }

    async fn rlopen(&self, fid: &Fid<Self::Fid>, flags: u32) -> Result<Fcall> {
        let realpath = {
            let realpath = fid.aux.realpath.read().await;
            realpath.clone()
        };

        let qid = get_qid(&realpath).await?;
        if !qid.typ.contains(QidType::DIR) {
            let oflags = nix::fcntl::OFlag::from_bits_truncate((flags & UNIX_FLAGS) as i32);
            let omode = nix::sys::stat::Mode::from_bits_truncate(0);
            let fd = nix::fcntl::open(&realpath, oflags, omode)?;

            {
                let mut file = fid.aux.file.lock().await;
                *file = Some(fs::File::from_std(unsafe { std::fs::File::from_raw_fd(fd) }));
            }
        }

        Ok(Fcall::Rlopen {
            qid: qid,
            iounit: 0,
        })
    }

    async fn rlcreate(
        &self,
        fid: &Fid<Self::Fid>,
        name: &str,
        flags: u32,
        mode: u32,
        _gid: u32,
    ) -> Result<Fcall> {
        let path = {
            let realpath = fid.aux.realpath.read().await;
            realpath.join(name)
        };
        let oflags = nix::fcntl::OFlag::from_bits_truncate((flags & UNIX_FLAGS) as i32);
        let omode = nix::sys::stat::Mode::from_bits_truncate(mode);
        let fd = nix::fcntl::open(&path, oflags, omode)?;

        let qid = get_qid(&path).await?;
        {
            let mut realpath = fid.aux.realpath.write().await;
            *realpath = path;
        }
        {
            let mut file = fid.aux.file.lock().await;
            *file = Some(fs::File::from_std(unsafe { std::fs::File::from_raw_fd(fd) }));
        }

        Ok(Fcall::Rlcreate {
            qid,
            iounit: 0,
        })
    }

    async fn rread(&self, fid: &Fid<Self::Fid>, offset: u64, count: u32) -> Result<Fcall> {
        let buf = {
            let mut file = fid.aux.file.lock().await;
            let file = file.as_mut().ok_or(INVALID_FID!())?;
            file.seek(SeekFrom::Start(offset)).await?;

            let mut buf = create_buffer(count as usize);
            let bytes = file.read(&mut buf[..]).await?;
            buf.truncate(bytes);
            buf
        };

        Ok(Fcall::Rread { data: Data(buf) })
    }

    async fn rwrite(&self, fid: &Fid<Self::Fid>, offset: u64, data: &Data) -> Result<Fcall> {
        let count = {
            let mut file = fid.aux.file.lock().await;
            let file = file.as_mut().ok_or(INVALID_FID!())?;
            file.seek(SeekFrom::Start(offset)).await?;
            file.write(&data.0).await? as u32
        };

        Ok(Fcall::Rwrite { count })
    }

    async fn rmkdir(
        &self,
        dfid: &Fid<Self::Fid>,
        _name: &str,
        _mode: u32,
        _gid: u32,
    ) -> Result<Fcall> {
        let path = {
            let realpath = dfid.aux.realpath.read().await;
            realpath.clone()
        };

        fs::create_dir(&path).await?;

        Ok(Fcall::Rmkdir {
            qid: get_qid(&path).await?,
        })
    }

    async fn rrenameat(
        &self,
        olddir: &Fid<Self::Fid>,
        oldname: &str,
        newdir: &Fid<Self::Fid>,
        newname: &str,
    ) -> Result<Fcall> {
        let oldpath = {
            let realpath = olddir.aux.realpath.read().await;
            realpath.join(oldname)
        };

        let newpath = {
            let realpath = newdir.aux.realpath.read().await;
            realpath.join(newname)
        };

        fs::rename(&oldpath, &newpath).await?;

        Ok(Fcall::Rrenameat)
    }

    async fn runlinkat(&self, dirfid: &Fid<Self::Fid>, name: &str, _flags: u32) -> Result<Fcall> {
        let path = {
            let realpath = dirfid.aux.realpath.read().await;
            realpath.join(name)
        };

        match fs::symlink_metadata(&path).await? {
            ref attr if attr.is_dir() => fs::remove_dir(&path).await?,
            _ => fs::remove_file(&path).await?,
        };

        Ok(Fcall::Runlinkat)
    }

    async fn rfsync(&self, fid: &Fid<Self::Fid>) -> Result<Fcall> {
        {
            let mut file = fid.aux.file.lock().await;
            file
                .as_mut()
                .ok_or(INVALID_FID!())?
                .sync_all().await?;
        }

        Ok(Fcall::Rfsync)
    }

    async fn rclunk(&self, _: &Fid<Self::Fid>) -> Result<Fcall> {
        Ok(Fcall::Rclunk)
    }

    async fn rstatfs(&self, fid: &Fid<Self::Fid>) -> Result<Fcall> {
        let path = {
            let realpath = fid.aux.realpath.read().await;
            realpath.clone()
        };

        //let fs = nix::sys::statvfs::statvfs(&path)?;
        let fs = tokio::task::spawn_blocking(move || nix::sys::statvfs::statvfs(&path)).await.unwrap()?;

        Ok(Fcall::Rstatfs {
            statfs: From::from(fs),
        })
    }
}

async fn unpfs_main(args: Vec<String>) -> rs9p::Result<i32> {
    if args.len() < 3 {
        eprintln!("Usage: {} proto!address!port mountpoint", args[0]);
        eprintln!("  where: proto = tcp | unix");
        return Ok(-1);
    }

    let (addr, mountpoint) = (&args[1], PathBuf::from(&args[2]));
    if !fs::metadata(&mountpoint).await?.is_dir() {
        return res!(io_err!(Other, "mount point must be a directory"));
    }

    println!("[*] Ready to accept clients: {}", addr);
    srv_async(
        Unpfs {
            realroot: mountpoint,
        },
        addr,
    ).await.and(Ok(0))
}

#[tokio::main]
async fn main() {
    env_logger::init();

    let args = std::env::args().collect();
    let exit_code = unpfs_main(args).await.unwrap_or_else(|e| {
        eprintln!("Error: {:?}", e);
        -1
    });

    std::process::exit(exit_code);
}
<commit_msg>unpfs: fix bug not joining the given directory name in rmkdir<commit_after>use {
    std::{
        io::SeekFrom,
        path::PathBuf,
        os::unix::{fs::PermissionsExt, io::FromRawFd},
    },
    filetime::FileTime,
    nix::libc::{O_CREAT, O_RDONLY, O_RDWR, O_TRUNC, O_WRONLY},
    async_trait::async_trait,
    tokio::{
        prelude::*,
        fs,
        stream::StreamExt,
        sync::{Mutex, RwLock},
    },
    rs9p::{*, srv::{srv_async, Fid, Filesystem}},
};

mod utils;
use crate::utils::*;

// Some clients will incorrectly set bits in 9p flags that don't make sense.
// For exmaple, the linux 9p kernel client propagates O_DIRECT to TCREATE and TOPEN
// and from there to the server.
// Processes on client machines set O_DIRECT to bypass the cache, but if
// the server uses O_DIRECT in the open or create, then subsequent server
// write and read system calls will fail, as O_DIRECT requires at minimum 512
// byte aligned data, and the data is almost always not aligned.
// While the linux kernel client is arguably broken, we won't be able
// to fix every kernel out there, and this is surely not the only buggy client
// we will see.
// The fix is to enumerate the set of flags we support and then and that with
// the flags received in a TCREATE or TOPEN. This nicely fixes a real problem
// we are seeing with a file system benchmark.
const UNIX_FLAGS: u32 = (O_WRONLY | O_RDONLY | O_RDWR | O_CREAT | O_TRUNC) as u32;

#[derive(Default)]
struct UnpfsFid {
    realpath: RwLock<PathBuf>,
    file: Mutex<Option<fs::File>>,
}

#[derive(Clone)]
struct Unpfs {
    realroot: PathBuf,
}

#[async_trait]
impl Filesystem for Unpfs {
    type Fid = UnpfsFid;

    async fn rattach(
        &self,
        fid: &Fid<Self::Fid>,
        _afid: Option<&Fid<Self::Fid>>,
        _uname: &str,
        _aname: &str,
        _n_uname: u32,
    ) -> Result<Fcall> {
        {
            let mut realpath = fid.aux.realpath.write().await;
            *realpath = PathBuf::from(&self.realroot);
        }

        Ok(Fcall::Rattach {
            qid: get_qid(&self.realroot).await?,
        })
    }

    async fn rwalk(
        &self,
        fid: &Fid<Self::Fid>,
        newfid: &Fid<Self::Fid>,
        wnames: &[String],
    ) -> Result<Fcall> {
        let mut wqids = Vec::new();
        let mut path = {
            let realpath = fid.aux.realpath.read().await;
            realpath.clone()
        };

        for (i, name) in wnames.iter().enumerate() {
            path.push(name);

            let qid = match get_qid(&path).await {
                Ok(qid) => qid,
                Err(e) => {
                    if i == 0 {
                        return Err(e);
                    } else {
                        break;
                    }
                }
            };

            wqids.push(qid);
        }

        {
            let mut new_realpath = newfid.aux.realpath.write().await;
            *new_realpath = path;
        }

        Ok(Fcall::Rwalk { wqids: wqids })
    }

    async fn rgetattr(&self, fid: &Fid<Self::Fid>, req_mask: GetattrMask) -> Result<Fcall> {
        let attr = {
            let realpath = fid.aux.realpath.read().await;
            fs::symlink_metadata(&*realpath).await?
        };

        Ok(Fcall::Rgetattr {
            valid: req_mask,
            qid: qid_from_attr(&attr),
            stat: From::from(attr),
        })
    }

    async fn rsetattr(
        &self,
        fid: &Fid<Self::Fid>,
        valid: SetattrMask,
        stat: &SetAttr,
    ) -> Result<Fcall> {
        let filepath = {
            let realpath = fid.aux.realpath.read().await;
            realpath.clone()
        };

        if valid.contains(SetattrMask::MODE) {
            fs::set_permissions(&filepath, PermissionsExt::from_mode(stat.mode)).await?;
        }

        if valid.intersects(SetattrMask::UID | SetattrMask::GID) {
            let uid = if valid.contains(SetattrMask::UID) {
                Some(nix::unistd::Uid::from_raw(stat.uid))
            } else {
                None
            };
            let gid = if valid.contains(SetattrMask::GID) {
                Some(nix::unistd::Gid::from_raw(stat.gid))
            } else {
                None
            };
            nix::unistd::chown(&filepath, uid, gid)?;
        }

        if valid.contains(SetattrMask::SIZE) {
            let _ = fs::File::open(&filepath).await?.set_len(stat.size).await?;
        }

        if valid.intersects(SetattrMask::ATIME_SET | SetattrMask::MTIME_SET) {
            let attr = fs::metadata(&filepath).await?;
            let atime = if valid.contains(SetattrMask::ATIME_SET) {
                FileTime::from_unix_time(stat.atime.sec as i64, stat.atime.nsec as u32)
            } else {
                FileTime::from_last_access_time(&attr)
            };

            let mtime = if valid.contains(SetattrMask::MTIME_SET) {
                FileTime::from_unix_time(stat.mtime.sec as i64, stat.mtime.nsec as u32)
            } else {
                FileTime::from_last_modification_time(&attr)
            };

            let _ = tokio::task::spawn_blocking(move || filetime::set_file_times(filepath, atime, mtime)).await;
        }

        Ok(Fcall::Rsetattr)
    }

    async fn rreadlink(&self, fid: &Fid<Self::Fid>) -> Result<Fcall> {
        let link = {
            let realpath = fid.aux.realpath.read().await;
            fs::read_link(&*realpath).await?
        };

        Ok(Fcall::Rreadlink {
            target: link.to_string_lossy().into_owned(),
        })
    }

    async fn rreaddir(&self, fid: &Fid<Self::Fid>, off: u64, count: u32) -> Result<Fcall> {
        let mut dirents = DirEntryData::new();

        let offset = if off == 0 {
            dirents.push(get_dirent_from(".", 0).await?);
            dirents.push(get_dirent_from("..", 1).await?);
            off
        } else {
            off - 1
        } as usize;

        let mut entries = {
            let realpath = fid.aux.realpath.read().await;
            fs::read_dir(&*realpath).await?.skip(offset)
        };

        let mut i = offset;
        while let Some(entry) = entries.next().await {
            let dirent = get_dirent(&entry?, 2 + i as u64).await?;
            if dirents.size() + dirent.size() > count {
                break;
            }
            dirents.push(dirent);
            i += 1;
        }

        Ok(Fcall::Rreaddir { data: dirents })
    }

    async fn rlopen(&self, fid: &Fid<Self::Fid>, flags: u32) -> Result<Fcall> {
        let realpath = {
            let realpath = fid.aux.realpath.read().await;
            realpath.clone()
        };

        let qid = get_qid(&realpath).await?;
        if !qid.typ.contains(QidType::DIR) {
            let oflags = nix::fcntl::OFlag::from_bits_truncate((flags & UNIX_FLAGS) as i32);
            let omode = nix::sys::stat::Mode::from_bits_truncate(0);
            let fd = nix::fcntl::open(&realpath, oflags, omode)?;

            {
                let mut file = fid.aux.file.lock().await;
                *file = Some(fs::File::from_std(unsafe { std::fs::File::from_raw_fd(fd) }));
            }
        }

        Ok(Fcall::Rlopen {
            qid: qid,
            iounit: 0,
        })
    }

    async fn rlcreate(
        &self,
        fid: &Fid<Self::Fid>,
        name: &str,
        flags: u32,
        mode: u32,
        _gid: u32,
    ) -> Result<Fcall> {
        let path = {
            let realpath = fid.aux.realpath.read().await;
            realpath.join(name)
        };
        let oflags = nix::fcntl::OFlag::from_bits_truncate((flags & UNIX_FLAGS) as i32);
        let omode = nix::sys::stat::Mode::from_bits_truncate(mode);
        let fd = nix::fcntl::open(&path, oflags, omode)?;

        let qid = get_qid(&path).await?;
        {
            let mut realpath = fid.aux.realpath.write().await;
            *realpath = path;
        }
        {
            let mut file = fid.aux.file.lock().await;
            *file = Some(fs::File::from_std(unsafe { std::fs::File::from_raw_fd(fd) }));
        }

        Ok(Fcall::Rlcreate {
            qid,
            iounit: 0,
        })
    }

    async fn rread(&self, fid: &Fid<Self::Fid>, offset: u64, count: u32) -> Result<Fcall> {
        let buf = {
            let mut file = fid.aux.file.lock().await;
            let file = file.as_mut().ok_or(INVALID_FID!())?;
            file.seek(SeekFrom::Start(offset)).await?;

            let mut buf = create_buffer(count as usize);
            let bytes = file.read(&mut buf[..]).await?;
            buf.truncate(bytes);
            buf
        };

        Ok(Fcall::Rread { data: Data(buf) })
    }

    async fn rwrite(&self, fid: &Fid<Self::Fid>, offset: u64, data: &Data) -> Result<Fcall> {
        let count = {
            let mut file = fid.aux.file.lock().await;
            let file = file.as_mut().ok_or(INVALID_FID!())?;
            file.seek(SeekFrom::Start(offset)).await?;
            file.write(&data.0).await? as u32
        };

        Ok(Fcall::Rwrite { count })
    }

    async fn rmkdir(
        &self,
        dfid: &Fid<Self::Fid>,
        name: &str,
        _mode: u32,
        _gid: u32,
    ) -> Result<Fcall> {
        let path = {
            let realpath = dfid.aux.realpath.read().await;
            realpath.join(name)
        };

        fs::create_dir(&path).await?;

        Ok(Fcall::Rmkdir {
            qid: get_qid(&path).await?,
        })
    }

    async fn rrenameat(
        &self,
        olddir: &Fid<Self::Fid>,
        oldname: &str,
        newdir: &Fid<Self::Fid>,
        newname: &str,
    ) -> Result<Fcall> {
        let oldpath = {
            let realpath = olddir.aux.realpath.read().await;
            realpath.join(oldname)
        };

        let newpath = {
            let realpath = newdir.aux.realpath.read().await;
            realpath.join(newname)
        };

        fs::rename(&oldpath, &newpath).await?;

        Ok(Fcall::Rrenameat)
    }

    async fn runlinkat(&self, dirfid: &Fid<Self::Fid>, name: &str, _flags: u32) -> Result<Fcall> {
        let path = {
            let realpath = dirfid.aux.realpath.read().await;
            realpath.join(name)
        };

        match fs::symlink_metadata(&path).await? {
            ref attr if attr.is_dir() => fs::remove_dir(&path).await?,
            _ => fs::remove_file(&path).await?,
        };

        Ok(Fcall::Runlinkat)
    }

    async fn rfsync(&self, fid: &Fid<Self::Fid>) -> Result<Fcall> {
        {
            let mut file = fid.aux.file.lock().await;
            file
                .as_mut()
                .ok_or(INVALID_FID!())?
                .sync_all().await?;
        }

        Ok(Fcall::Rfsync)
    }

    async fn rclunk(&self, _: &Fid<Self::Fid>) -> Result<Fcall> {
        Ok(Fcall::Rclunk)
    }

    async fn rstatfs(&self, fid: &Fid<Self::Fid>) -> Result<Fcall> {
        let path = {
            let realpath = fid.aux.realpath.read().await;
            realpath.clone()
        };

        //let fs = nix::sys::statvfs::statvfs(&path)?;
        let fs = tokio::task::spawn_blocking(move || nix::sys::statvfs::statvfs(&path)).await.unwrap()?;

        Ok(Fcall::Rstatfs {
            statfs: From::from(fs),
        })
    }
}

async fn unpfs_main(args: Vec<String>) -> rs9p::Result<i32> {
    if args.len() < 3 {
        eprintln!("Usage: {} proto!address!port mountpoint", args[0]);
        eprintln!("  where: proto = tcp | unix");
        return Ok(-1);
    }

    let (addr, mountpoint) = (&args[1], PathBuf::from(&args[2]));
    if !fs::metadata(&mountpoint).await?.is_dir() {
        return res!(io_err!(Other, "mount point must be a directory"));
    }

    println!("[*] Ready to accept clients: {}", addr);
    srv_async(
        Unpfs {
            realroot: mountpoint,
        },
        addr,
    ).await.and(Ok(0))
}

#[tokio::main]
async fn main() {
    env_logger::init();

    let args = std::env::args().collect();
    let exit_code = unpfs_main(args).await.unwrap_or_else(|e| {
        eprintln!("Error: {:?}", e);
        -1
    });

    std::process::exit(exit_code);
}
<|endoftext|>


<commit_before>extern crate timely;
extern crate differential_dataflow;

use std::sync::{Arc, Mutex};
use std::net::TcpListener;
use std::time::Duration;

use timely::dataflow::operators::{Map, Inspect};
use timely::progress::nested::product::Product;
use timely::progress::timestamp::RootTimestamp;
use timely::logging::TimelyEvent;
use timely::communication::logging::{CommunicationEvent, CommunicationSetup};
use timely::dataflow::operators::Filter;
use timely::dataflow::operators::capture::{EventReader, Replay};

use differential_dataflow::AsCollection;
use differential_dataflow::operators::{Consolidate, Join};
// use differential_dataflow::logging::DifferentialEvent;

fn main() {

    let mut args = ::std::env::args();
    args.next().unwrap();

    let work_peers = args.next().expect("Must provide number of source peers").parse::<usize>().expect("Source peers must be an unsigned integer");
    let comm_peers = args.next().expect("Must provide number of source peers").parse::<usize>().expect("Source peers must be an unsigned integer");

    let t_listener = TcpListener::bind("127.0.0.1:8000").unwrap();
    let d_listener = TcpListener::bind("127.0.0.1:9000").unwrap();
    let t_sockets =
    Arc::new(Mutex::new((0..work_peers).map(|_| {
            let socket = t_listener.incoming().next().unwrap().unwrap();
            socket.set_nonblocking(true).expect("failed to set nonblocking");
            Some(socket)
        }).collect::<Vec<_>>()));
    let d_sockets =
    Arc::new(Mutex::new((0..comm_peers).map(|_| {
            let socket = d_listener.incoming().next().unwrap().unwrap();
            socket.set_nonblocking(true).expect("failed to set nonblocking");
            Some(socket)
        }).collect::<Vec<_>>()));

    timely::execute_from_args(std::env::args(), move |worker| {

        let index = worker.index();
        let peers = worker.peers();

        let t_streams =
        t_sockets
            .lock()
            .unwrap()
            .iter_mut()
            .enumerate()
            .filter(|(i, _)| *i % peers == index)
            .map(move |(_, s)| s.take().unwrap())
            .map(|r| EventReader::<Product<RootTimestamp, Duration>, (Duration, usize, TimelyEvent),_>::new(r))
            .collect::<Vec<_>>();

        let d_streams =
        d_sockets
            .lock()
            .unwrap()
            .iter_mut()
            .enumerate()
            .filter(|(i, _)| *i % peers == index)
            .map(move |(_, s)| s.take().unwrap())
            .map(|r| EventReader::<Product<RootTimestamp, Duration>, (Duration, CommunicationSetup, CommunicationEvent),_>::new(r))
            .collect::<Vec<_>>();

        worker.dataflow::<_,_,_>(|scope| {

            let t_events = t_streams.replay_into(scope);
            let d_events = d_streams.replay_into(scope);

            // let operates =
            t_events
                .filter(|x| x.1 == 0)
                .flat_map(move |(ts, _worker, datum)| {
                    let ts = Duration::from_secs(ts.as_secs() + 1);
                    if let TimelyEvent::Channels(event) = datum {
                        Some((event, RootTimestamp::new(ts), 1))
                    }
                    else { None }
                })
                .as_collection()
                // .consolidate()
                .inspect(|x| println!("CHANNEL\t{:?}", x));
                ;

            t_events
                .filter(|x| x.1 == 0)
                .flat_map(move |(ts, _worker, datum)| {
                    let ts = Duration::from_secs(ts.as_secs() + 1);
                    if let TimelyEvent::CommChannels(event) = datum {
                        Some((event, RootTimestamp::new(ts), 1))
                    }
                    else { None }
                })
                .as_collection()
                // .consolidate()
                .inspect(|x| println!("COMM_CHANNEL\t{:?}", x));
                ;

            // let memory =
            d_events
                // .inspect(|x| println!("MESSAGE\t{:?}", x))
                .flat_map(|(ts, _worker, datum)| {
                    let ts = Duration::from_secs(ts.as_secs() + 1);
                    if let CommunicationEvent::Message(x) = datum {
                        Some((x.header.channel, RootTimestamp::new(ts), x.header.length as isize))
                    }
                    else { None }
                })
                .as_collection()
                .consolidate()
                .inspect(|x| println!("MESSAGE\t{:?}", x));
                ;

            // operates
            //     .inspect(|x| println!("OPERATES: {:?}", x))
            //     .semijoin(&memory)
            //     .inspect(|x| println!("{:?}", x));

        });

    }).unwrap(); // asserts error-free execution
}
<commit_msg>made functional<commit_after>extern crate timely;
extern crate differential_dataflow;

use std::sync::{Arc, Mutex};
use std::net::TcpListener;
use std::time::Duration;

use timely::dataflow::operators::{Map, Inspect};
use timely::progress::nested::product::Product;
use timely::progress::timestamp::RootTimestamp;
use timely::logging::TimelyEvent;
use timely::communication::logging::{CommunicationEvent, CommunicationSetup};
use timely::dataflow::operators::Filter;
use timely::dataflow::operators::capture::{EventReader, Replay};

use differential_dataflow::AsCollection;
use differential_dataflow::operators::{Consolidate, Join};
// use differential_dataflow::logging::DifferentialEvent;

fn main() {

    let mut args = ::std::env::args();
    args.next().unwrap();

    let work_peers: usize = args.next().expect("Must provide number of source peers").parse().expect("Source peers must be an unsigned integer");
    let comm_peers: usize = args.next().expect("Must provide number of source peers").parse().expect("Comms peers must be an unsigned integer");
    let granularity: u64  = args.next().expect("Must provide report granularity in seconds").parse().expect("Granularity must be an unsigned integer");

	println!("starting with work peers: {}, comm peers: {}, granularity: {}", work_peers, comm_peers, granularity);

    let t_listener = TcpListener::bind("127.0.0.1:8000").unwrap();
    let d_listener = TcpListener::bind("127.0.0.1:9000").unwrap();
    let t_sockets =
    Arc::new(Mutex::new((0..work_peers).map(|_| {
            let socket = t_listener.incoming().next().unwrap().unwrap();
            socket.set_nonblocking(true).expect("failed to set nonblocking");
            Some(socket)
        }).collect::<Vec<_>>()));
    let d_sockets =
    Arc::new(Mutex::new((0..comm_peers).map(|_| {
            let socket = d_listener.incoming().next().unwrap().unwrap();
            socket.set_nonblocking(true).expect("failed to set nonblocking");
            Some(socket)
        }).collect::<Vec<_>>()));

    timely::execute_from_args(std::env::args(), move |worker| {

        let index = worker.index();
        let peers = worker.peers();

        let t_streams =
        t_sockets
            .lock()
            .unwrap()
            .iter_mut()
            .enumerate()
            .filter(|(i, _)| *i % peers == index)
            .map(move |(_, s)| s.take().unwrap())
            .map(|r| EventReader::<Product<RootTimestamp, Duration>, (Duration, usize, TimelyEvent),_>::new(r))
            .collect::<Vec<_>>();

        let d_streams =
        d_sockets
            .lock()
            .unwrap()
            .iter_mut()
            .enumerate()
            .filter(|(i, _)| *i % peers == index)
            .map(move |(_, s)| s.take().unwrap())
            .map(|r| EventReader::<Product<RootTimestamp, Duration>, (Duration, CommunicationSetup, CommunicationEvent),_>::new(r))
            .collect::<Vec<_>>();

        worker.dataflow::<_,_,_>(|scope| {

            let t_events = t_streams.replay_into(scope);
            let d_events = d_streams.replay_into(scope);

            // let operates =
            t_events
                .filter(|x| x.1 == 0)
                .flat_map(move |(ts, _worker, datum)| {
                    let ts = Duration::from_secs((ts.as_secs()/granularity + 1) * granularity);
                    if let TimelyEvent::Channels(event) = datum {
                        Some((event, RootTimestamp::new(ts), 1))
                    }
                    else { None }
                })
//                .inspect(|x| println!("???\t{:?}", x))
                .as_collection()
                // .consolidate()
                .inspect(|x| println!("CHANNEL\t{:?}", x));
                ;

            t_events
                .filter(|x| x.1 == 0)
                .flat_map(move |(ts, _worker, datum)| {
                    let ts = Duration::from_secs((ts.as_secs()/granularity + 1) * granularity);
                    if let TimelyEvent::CommChannels(event) = datum {
                        Some((event, RootTimestamp::new(ts), 1))
                    }
                    else { None }
                })
                .as_collection()
                // .consolidate()
                .inspect(|x| println!("COMM_CHANNEL\t{:?}", x));
                ;

            // let memory =
            d_events
                // .inspect(|x| println!("MESSAGE\t{:?}", x))
                .flat_map(move |(ts, _worker, datum)| {
                    let ts = Duration::from_secs((ts.as_secs()/granularity + 1) * granularity);
                    if let CommunicationEvent::Message(x) = datum {
                        Some((x.header.channel, RootTimestamp::new(ts), x.header.length as isize))
                    }
                    else { None }
                })
                .as_collection()
                .consolidate()
                .inspect(|x| println!("MESSAGE\t{:?}", x));
                ;

            // operates
            //     .inspect(|x| println!("OPERATES: {:?}", x))
            //     .semijoin(&memory)
            //     .inspect(|x| println!("{:?}", x));

        });

    }).unwrap(); // asserts error-free execution
}
<|endoftext|>