<commit_before>use std::ops::{Deref, DerefMut};
use std::fmt;

use std::convert::From;

use combine::primitives::SourcePosition;

/// Struct representing a position within a source code file.
///
/// This represents positions using `i32`s because that's how
/// positions are represented in `combine` (the parsing library
/// that we will use for the Mnemosyne parser). I personally would
/// have used `usize`s...
#[derive(Copy, Clone, PartialEq, Debug)]
pub struct Position {
    pub col: i32,
    pub row: i32,
    pub raw: i32
}

impl Position {

    /// Create a new `Position `at the given column and row.
    #[inline]
    pub fn new(col: i32, row: i32) -> Self {
        Position { col: col
                 , row: row
                 , raw: col + row
                 }
    }

}

impl From<SourcePosition> for Position {
    /// Create a new `Position` from a `combine` `SourcePosition`.
    ///
    /// # Example
    /// ```
    /// let sp = SourcePosition { column: 1, line: 1 };
    /// assert_eq!(Position::from(sp), Position::new(1,1));
    /// ```
    fn from(p: SourcePosition) -> Self { Position::new(p.column, p.line) }
}

impl From<(i32,i32)> for Position {
    /// Create a new `Position` from a tuple of i32s.
    ///
    /// # Example
    /// ```
    /// let tuple: (i32,i32) = (1,1);
    /// assert_eq!(Position::from(tuple), Position::new(1,1));
    /// ```
    fn from((col, row): (i32,i32)) -> Self { Position::new(col,row) }
}

impl fmt::Display for Position {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "line {}, column {}", self.row, self.col)
    }
}

#[derive(Clone, Debug)]
pub struct Positional<T> {
    pub pos: Position,
    pub value: T
}

impl<T> Positional<T> {
    /// Create a new Positional marker at the given position.
    pub fn at(col: i32, row: i32, value: T) -> Positional<T> {
        Positional { pos: Position::new(col, row)
                   , value: value }
    }

    pub fn value(&self) -> &T { &self.value }
}

impl<T> fmt::Display for Positional<T>
where T: fmt::Display {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} at {}", self.value, self.pos)
    }
}

impl<T> PartialEq for Positional<T>
where T: PartialEq {
    fn eq(&self, other: &Positional<T>) -> bool {
        self.value == other.value
    }
}

impl<T> Deref for Positional<T> {
    type Target = T;
    fn deref(&self) -> &T {
        &self.value
    }
}

impl<T> DerefMut for Positional<T> {
    fn deref_mut(&mut self) -> &mut T {
        &mut self.value
    }
}
<commit_msg>Derive more traits on Position<commit_after>use std::ops::{Deref, DerefMut};
use std::fmt;

use std::convert::From;

use combine::primitives::SourcePosition;

/// Struct representing a position within a source code file.
///
/// This represents positions using `i32`s because that's how
/// positions are represented in `combine` (the parsing library
/// that we will use for the Mnemosyne parser). I personally would
/// have used `usize`s...
#[derive(Copy, Clone, PartialEq, Eq, Debug, PartialOrd, Ord)]
pub struct Position {
    pub col: i32,
    pub row: i32,
    pub raw: i32
}

impl Position {

    /// Create a new `Position `at the given column and row.
    #[inline]
    pub fn new(col: i32, row: i32) -> Self {
        Position { col: col
                 , row: row
                 , raw: col + row
                 }
    }

}

impl From<SourcePosition> for Position {
    /// Create a new `Position` from a `combine` `SourcePosition`.
    ///
    /// # Example
    /// ```
    /// let sp = SourcePosition { column: 1, line: 1 };
    /// assert_eq!(Position::from(sp), Position::new(1,1));
    /// ```
    fn from(p: SourcePosition) -> Self { Position::new(p.column, p.line) }
}

impl From<(i32,i32)> for Position {
    /// Create a new `Position` from a tuple of i32s.
    ///
    /// # Example
    /// ```
    /// let tuple: (i32,i32) = (1,1);
    /// assert_eq!(Position::from(tuple), Position::new(1,1));
    /// ```
    fn from((col, row): (i32,i32)) -> Self { Position::new(col,row) }
}

impl fmt::Display for Position {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "line {}, column {}", self.row, self.col)
    }
}

#[derive(Clone, Debug)]
pub struct Positional<T> {
    pub pos: Position,
    pub value: T
}

impl<T> Positional<T> {
    /// Create a new Positional marker at the given position.
    pub fn at(col: i32, row: i32, value: T) -> Positional<T> {
        Positional { pos: Position::new(col, row)
                   , value: value }
    }

    pub fn value(&self) -> &T { &self.value }
}

impl<T> fmt::Display for Positional<T>
where T: fmt::Display {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} at {}", self.value, self.pos)
    }
}

impl<T> PartialEq for Positional<T>
where T: PartialEq {
    fn eq(&self, other: &Positional<T>) -> bool {
        self.value == other.value
    }
}

impl<T> Deref for Positional<T> {
    type Target = T;
    fn deref(&self) -> &T {
        &self.value
    }
}

impl<T> DerefMut for Positional<T> {
    fn deref_mut(&mut self) -> &mut T {
        &mut self.value
    }
}
<|endoftext|>


<commit_before>#[repr(C)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum DataType {
    TBIT,
    TBYTE,
    TSBYTE,
    TLOGICAL,
    TSTRING,
    TUSHORT,
    TSHORT,
    TUINT,
    TINT,
    TULONG,
    TLONG,
    TLONGLONG,
    TFLOAT,
    TDOUBLE,
    TCOMPLEX,
    TDBLCOMPLEX,
}

macro_rules! datatype_into_impl {
    ($t: ty) => (
        impl From<DataType> for $t {
            fn from(original: DataType) -> $t {
                match original {
                    DataType::TBIT => 1,
                    DataType::TBYTE => 11,
                    DataType::TSBYTE => 12,
                    DataType::TLOGICAL => 14,
                    DataType::TSTRING => 16,
                    DataType::TUSHORT => 20,
                    DataType::TSHORT => 21,
                    DataType::TUINT => 30,
                    DataType::TINT => 31,
                    DataType::TULONG => 40,
                    DataType::TLONG => 41,
                    DataType::TLONGLONG => 81,
                    DataType::TFLOAT => 42,
                    DataType::TDOUBLE => 82,
                    DataType::TCOMPLEX => 83,
                    DataType::TDBLCOMPLEX => 163,
                }
            }
        }
    )
}

datatype_into_impl!(u8);
datatype_into_impl!(i32);
datatype_into_impl!(u32);
datatype_into_impl!(i64);
datatype_into_impl!(u64);

#[repr(C)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum ImageType {
    BYTE_IMG,
    SHORT_IMG,
    LONG_IMG,
    LONGLONG_IMG,
    FLOAT_IMG,
    DOUBLE_IMG,
}

macro_rules! imagetype_into_impl {
    ($t: ty) => (
        impl From<ImageType> for $t {
            fn from(original: ImageType) -> $t {
                match original {
                    ImageType::BYTE_IMG => 8,
                    ImageType::SHORT_IMG => 16,
                    ImageType::LONG_IMG => 32,
                    ImageType::LONGLONG_IMG => 64,
                    ImageType::FLOAT_IMG => -32,
                    ImageType::DOUBLE_IMG => -64,
                }
            }
        }
        )
}

imagetype_into_impl!(i8);
imagetype_into_impl!(i32);
imagetype_into_impl!(i64);

#[repr(C)]
#[derive(Debug, PartialEq, Eq)]
pub enum HduType {
    IMAGE_HDU,
    ASCII_TBL,
    BINARY_TBL,
    ANY_HDU,
}

macro_rules! hdutype_into_impl {
    ($t: ty) => (
        impl From<HduType> for $t {
            fn from(original: HduType) -> $t {
                match original {
                    HduType::IMAGE_HDU => 0,
                    HduType::ASCII_TBL => 1,
                    HduType::BINARY_TBL => 2,
                    HduType::ANY_HDU => -1,
                }
            }
        }
    )
}

hdutype_into_impl!(i8);
hdutype_into_impl!(i32);
hdutype_into_impl!(i64);

#[repr(C)]
#[derive(Debug, PartialEq, Eq)]
pub enum FileOpenMode {
    READONLY,
    READWRITE,
}

macro_rules! fileopenmode_into_impl {
    ($t: ty) => (
        impl From<FileOpenMode> for $t {
            fn from(original: FileOpenMode) -> $t {
                match original {
                    FileOpenMode::READONLY => 0,
                    FileOpenMode::READWRITE => 1,
                }
            }
        }
        )
}

fileopenmode_into_impl!(u8);
fileopenmode_into_impl!(u32);
fileopenmode_into_impl!(u64);
fileopenmode_into_impl!(i8);
fileopenmode_into_impl!(i32);
fileopenmode_into_impl!(i64);

#[repr(C)]
#[derive(Debug, PartialEq, Eq)]
pub enum CaseSensitivity {
    CASEINSEN,
    CASESEN,
}

macro_rules! casesensitivity_into_impl {
    ($t: ty) => (
        impl From<CaseSensitivity> for $t {
            fn from(original: CaseSensitivity) -> $t {
                match original {
                    CaseSensitivity::CASEINSEN => 0,
                    CaseSensitivity::CASESEN => 1,
                }
            }
        }
        )
}

casesensitivity_into_impl!(u8);
casesensitivity_into_impl!(u32);
casesensitivity_into_impl!(u64);
casesensitivity_into_impl!(i8);
casesensitivity_into_impl!(i32);
casesensitivity_into_impl!(i64);

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

    #[test]
    fn image_types() {
        assert_eq!(i8::from(ImageType::BYTE_IMG), 8);
        assert_eq!(i8::from(ImageType::SHORT_IMG), 16);
        assert_eq!(i8::from(ImageType::LONG_IMG), 32);
        assert_eq!(i8::from(ImageType::LONGLONG_IMG), 64);
        assert_eq!(i8::from(ImageType::FLOAT_IMG), -32);
        assert_eq!(i8::from(ImageType::DOUBLE_IMG), -64);
    }

    #[test]
    fn hdu_types() {
        assert_eq!(i32::from(HduType::IMAGE_HDU), 0);
        assert_eq!(i32::from(HduType::ASCII_TBL), 1);
        assert_eq!(i32::from(HduType::BINARY_TBL), 2);
        assert_eq!(i32::from(HduType::ANY_HDU), -1);
    }

    #[test]
    fn file_open_modes() {
        assert_eq!(u8::from(FileOpenMode::READONLY), 0);
        assert_eq!(u8::from(FileOpenMode::READWRITE), 1);
    }

    #[test]
    fn case_sensitivity() {
        assert_eq!(u8::from(CaseSensitivity::CASESEN), 1);
        assert_eq!(u8::from(CaseSensitivity::CASEINSEN), 0);
    }

    #[test]
    fn converting_from_data_type() {
        assert_eq!(u8::from(DataType::TBIT), 1);
        assert_eq!(u8::from(DataType::TBYTE), 11);
        assert_eq!(u8::from(DataType::TLOGICAL), 14);
        assert_eq!(u8::from(DataType::TSTRING), 16);
        assert_eq!(u8::from(DataType::TSHORT), 21);
        assert_eq!(u8::from(DataType::TLONG), 41);
        assert_eq!(u8::from(DataType::TFLOAT), 42);
        assert_eq!(u8::from(DataType::TDOUBLE), 82);
    }

}
<commit_msg>Prevent non-camel-case type warnings<commit_after>#[repr(C)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum DataType {
    TBIT,
    TBYTE,
    TSBYTE,
    TLOGICAL,
    TSTRING,
    TUSHORT,
    TSHORT,
    TUINT,
    TINT,
    TULONG,
    TLONG,
    TLONGLONG,
    TFLOAT,
    TDOUBLE,
    TCOMPLEX,
    TDBLCOMPLEX,
}

macro_rules! datatype_into_impl {
    ($t: ty) => (
        impl From<DataType> for $t {
            fn from(original: DataType) -> $t {
                match original {
                    DataType::TBIT => 1,
                    DataType::TBYTE => 11,
                    DataType::TSBYTE => 12,
                    DataType::TLOGICAL => 14,
                    DataType::TSTRING => 16,
                    DataType::TUSHORT => 20,
                    DataType::TSHORT => 21,
                    DataType::TUINT => 30,
                    DataType::TINT => 31,
                    DataType::TULONG => 40,
                    DataType::TLONG => 41,
                    DataType::TLONGLONG => 81,
                    DataType::TFLOAT => 42,
                    DataType::TDOUBLE => 82,
                    DataType::TCOMPLEX => 83,
                    DataType::TDBLCOMPLEX => 163,
                }
            }
        }
    )
}

datatype_into_impl!(u8);
datatype_into_impl!(i32);
datatype_into_impl!(u32);
datatype_into_impl!(i64);
datatype_into_impl!(u64);

#[allow(non_camel_case_types)]
#[repr(C)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum ImageType {
    BYTE_IMG,
    SHORT_IMG,
    LONG_IMG,
    LONGLONG_IMG,
    FLOAT_IMG,
    DOUBLE_IMG,
}

macro_rules! imagetype_into_impl {
    ($t: ty) => (
        impl From<ImageType> for $t {
            fn from(original: ImageType) -> $t {
                match original {
                    ImageType::BYTE_IMG => 8,
                    ImageType::SHORT_IMG => 16,
                    ImageType::LONG_IMG => 32,
                    ImageType::LONGLONG_IMG => 64,
                    ImageType::FLOAT_IMG => -32,
                    ImageType::DOUBLE_IMG => -64,
                }
            }
        }
        )
}

imagetype_into_impl!(i8);
imagetype_into_impl!(i32);
imagetype_into_impl!(i64);

#[allow(non_camel_case_types)]
#[repr(C)]
#[derive(Debug, PartialEq, Eq)]
pub enum HduType {
    IMAGE_HDU,
    ASCII_TBL,
    BINARY_TBL,
    ANY_HDU,
}

macro_rules! hdutype_into_impl {
    ($t: ty) => (
        impl From<HduType> for $t {
            fn from(original: HduType) -> $t {
                match original {
                    HduType::IMAGE_HDU => 0,
                    HduType::ASCII_TBL => 1,
                    HduType::BINARY_TBL => 2,
                    HduType::ANY_HDU => -1,
                }
            }
        }
    )
}

hdutype_into_impl!(i8);
hdutype_into_impl!(i32);
hdutype_into_impl!(i64);

#[repr(C)]
#[derive(Debug, PartialEq, Eq)]
pub enum FileOpenMode {
    READONLY,
    READWRITE,
}

macro_rules! fileopenmode_into_impl {
    ($t: ty) => (
        impl From<FileOpenMode> for $t {
            fn from(original: FileOpenMode) -> $t {
                match original {
                    FileOpenMode::READONLY => 0,
                    FileOpenMode::READWRITE => 1,
                }
            }
        }
        )
}

fileopenmode_into_impl!(u8);
fileopenmode_into_impl!(u32);
fileopenmode_into_impl!(u64);
fileopenmode_into_impl!(i8);
fileopenmode_into_impl!(i32);
fileopenmode_into_impl!(i64);

#[repr(C)]
#[derive(Debug, PartialEq, Eq)]
pub enum CaseSensitivity {
    CASEINSEN,
    CASESEN,
}

macro_rules! casesensitivity_into_impl {
    ($t: ty) => (
        impl From<CaseSensitivity> for $t {
            fn from(original: CaseSensitivity) -> $t {
                match original {
                    CaseSensitivity::CASEINSEN => 0,
                    CaseSensitivity::CASESEN => 1,
                }
            }
        }
        )
}

casesensitivity_into_impl!(u8);
casesensitivity_into_impl!(u32);
casesensitivity_into_impl!(u64);
casesensitivity_into_impl!(i8);
casesensitivity_into_impl!(i32);
casesensitivity_into_impl!(i64);

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

    #[test]
    fn image_types() {
        assert_eq!(i8::from(ImageType::BYTE_IMG), 8);
        assert_eq!(i8::from(ImageType::SHORT_IMG), 16);
        assert_eq!(i8::from(ImageType::LONG_IMG), 32);
        assert_eq!(i8::from(ImageType::LONGLONG_IMG), 64);
        assert_eq!(i8::from(ImageType::FLOAT_IMG), -32);
        assert_eq!(i8::from(ImageType::DOUBLE_IMG), -64);
    }

    #[test]
    fn hdu_types() {
        assert_eq!(i32::from(HduType::IMAGE_HDU), 0);
        assert_eq!(i32::from(HduType::ASCII_TBL), 1);
        assert_eq!(i32::from(HduType::BINARY_TBL), 2);
        assert_eq!(i32::from(HduType::ANY_HDU), -1);
    }

    #[test]
    fn file_open_modes() {
        assert_eq!(u8::from(FileOpenMode::READONLY), 0);
        assert_eq!(u8::from(FileOpenMode::READWRITE), 1);
    }

    #[test]
    fn case_sensitivity() {
        assert_eq!(u8::from(CaseSensitivity::CASESEN), 1);
        assert_eq!(u8::from(CaseSensitivity::CASEINSEN), 0);
    }

    #[test]
    fn converting_from_data_type() {
        assert_eq!(u8::from(DataType::TBIT), 1);
        assert_eq!(u8::from(DataType::TBYTE), 11);
        assert_eq!(u8::from(DataType::TLOGICAL), 14);
        assert_eq!(u8::from(DataType::TSTRING), 16);
        assert_eq!(u8::from(DataType::TSHORT), 21);
        assert_eq!(u8::from(DataType::TLONG), 41);
        assert_eq!(u8::from(DataType::TFLOAT), 42);
        assert_eq!(u8::from(DataType::TDOUBLE), 82);
    }

}
<|endoftext|>


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

use bincode::serde::DeserializeError;
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
use std::collections::hash_map::HashMap;
use std::cell::{RefCell};
use std::io::{Error, ErrorKind};
use std::slice;
use std::fmt::{self, Debug, Formatter};
use std::cmp::{PartialEq};
use std::ops::Deref;
use std::mem;
use std::usize;

use uuid::Uuid;

struct ServerRecord {
    sender: OsIpcSender,
    conn_sender: mpsc::Sender<bool>,
    conn_receiver: Mutex<mpsc::Receiver<bool>>,
}

impl ServerRecord {
    fn new(sender: OsIpcSender) -> ServerRecord {
        let (tx, rx) = mpsc::channel::<bool>();
        ServerRecord {
            sender: sender,
            conn_sender: tx,
            conn_receiver: Mutex::new(rx),
        }
    }

    fn accept(&self) {
        self.conn_receiver.lock().unwrap().recv().unwrap();
    }

    fn connect(&self) {
        self.conn_sender.send(true).unwrap();
    }
}

lazy_static! {
    static ref ONE_SHOT_SERVERS: Mutex<HashMap<String,ServerRecord>> = Mutex::new(HashMap::new());
}

struct MpscChannelMessage(Vec<u8>, Vec<OsIpcChannel>, Vec<OsIpcSharedMemory>);

pub fn channel() -> Result<(OsIpcSender, OsIpcReceiver),MpscError> {
    let (base_sender, base_receiver) = mpsc::channel::<MpscChannelMessage>();
    Ok((OsIpcSender::new(base_sender), OsIpcReceiver::new(base_receiver)))
}

pub struct OsIpcReceiver {
    receiver: RefCell<Option<mpsc::Receiver<MpscChannelMessage>>>,
}

impl PartialEq for OsIpcReceiver {
    fn eq(&self, other: &OsIpcReceiver) -> bool {
        self.receiver.borrow().as_ref().map(|rx| rx as *const _) ==
            other.receiver.borrow().as_ref().map(|rx| rx as *const _)
    }
}

// Can't derive, as mpsc::Receiver doesn't implement Debug.
impl fmt::Debug for OsIpcReceiver {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // Not sure there is anything useful we could print here.
        write!(f, "OsIpcReceiver {{ .. }}")
    }
}

impl OsIpcReceiver {
    fn new(receiver: mpsc::Receiver<MpscChannelMessage>) -> OsIpcReceiver {
        OsIpcReceiver {
            receiver: RefCell::new(Some(receiver)),
        }
    }

    pub fn consume(&self) -> OsIpcReceiver {
        let receiver = self.receiver.borrow_mut().take();
        OsIpcReceiver::new(receiver.unwrap())
    }

    pub fn recv(&self) -> Result<(Vec<u8>, Vec<OsOpaqueIpcChannel>, Vec<OsIpcSharedMemory>),MpscError> {
        let r = self.receiver.borrow();
        match r.as_ref().unwrap().recv() {
            Ok(MpscChannelMessage(d,c,s)) => Ok((d,
                                                 c.into_iter().map(OsOpaqueIpcChannel::new).collect(),
                                                 s)),
            Err(_) => Err(MpscError::ChannelClosedError),
        }
    }

    pub fn try_recv(&self) -> Result<(Vec<u8>, Vec<OsOpaqueIpcChannel>, Vec<OsIpcSharedMemory>),MpscError> {
        let r = self.receiver.borrow();
        match r.as_ref().unwrap().try_recv() {
            Ok(MpscChannelMessage(d,c,s)) => Ok((d,
                                                 c.into_iter().map(OsOpaqueIpcChannel::new).collect(),
                                                 s)),
            Err(_) => Err(MpscError::ChannelClosedError),
        }
    }
}

unsafe impl Send for OsIpcReceiver { }
unsafe impl Sync for OsIpcReceiver { }

#[derive(Clone)]
pub struct OsIpcSender {
    sender: RefCell<mpsc::Sender<MpscChannelMessage>>,
}

impl PartialEq for OsIpcSender {
    fn eq(&self, other: &OsIpcSender) -> bool {
        &*self.sender.borrow() as *const _ ==
            &*other.sender.borrow() as *const _
    }
}

// Can't derive, as mpsc::Sender doesn't implement Debug.
impl fmt::Debug for OsIpcSender {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // Not sure there is anything useful we could print here.
        write!(f, "OsIpcSender {{ .. }}")
    }
}

unsafe impl Send for OsIpcSender { }
unsafe impl Sync for OsIpcSender { }

impl OsIpcSender {
    fn new(sender: mpsc::Sender<MpscChannelMessage>) -> OsIpcSender {
        OsIpcSender {
            sender: RefCell::new(sender),
        }
    }

    pub fn connect(name: String) -> Result<OsIpcSender,MpscError> {
        let record = ONE_SHOT_SERVERS.lock().unwrap().remove(&name).unwrap();
        record.connect();
        Ok(record.sender)
    }

    pub fn get_max_fragment_size() -> usize {
        usize::MAX
    }

    pub fn send(&self,
                data: &[u8],
                ports: Vec<OsIpcChannel>,
                shared_memory_regions: Vec<OsIpcSharedMemory>)
                -> Result<(),MpscError>
    {
        match self.sender.borrow().send(MpscChannelMessage(data.to_vec(), ports, shared_memory_regions)) {
            Err(_) => Err(MpscError::ChannelClosedError),
            Ok(_) => Ok(()),
        }
    }
}

pub struct OsIpcReceiverSet {
    last_index: usize,
    receiver_ids: Vec<usize>,
    receivers: Vec<OsIpcReceiver>,
}

impl OsIpcReceiverSet {
    pub fn new() -> Result<OsIpcReceiverSet,MpscError> {
        Ok(OsIpcReceiverSet {
            last_index: 0,
            receiver_ids: vec![],
            receivers: vec![],
        })
    }

    pub fn add(&mut self, receiver: OsIpcReceiver) -> Result<i64,MpscError> {
        self.last_index += 1;
        self.receiver_ids.push(self.last_index);
        self.receivers.push(receiver.consume());
        Ok(self.last_index as i64)
    }

    pub fn select(&mut self) -> Result<Vec<OsIpcSelectionResult>,MpscError> {
        let mut receivers: Vec<Option<mpsc::Receiver<MpscChannelMessage>>> = Vec::with_capacity(self.receivers.len());
        let mut r_id: i64 = -1;
        let mut r_index: usize = 0;

        {
            let select = mpsc::Select::new();
            // we *must* allocate exact capacity for this, because the Handles *can't move*
            let mut handles: Vec<mpsc::Handle<MpscChannelMessage>> = Vec::with_capacity(self.receivers.len());

            for r in &self.receivers {
                let inner_r = mem::replace(&mut *r.receiver.borrow_mut(), None);
                receivers.push(inner_r);
            }

            for r in &receivers {
                unsafe {
                    handles.push(select.handle(r.as_ref().unwrap()));
                    handles.last_mut().unwrap().add();
                }
            }

            let id = select.wait();

            for (index,h) in handles.iter().enumerate() {
                if h.id() == id {
                    r_index = index;
                    r_id = self.receiver_ids[index] as i64;
                    break;
                }
            }
        }

        // put the receivers back
        for (index,r) in self.receivers.iter().enumerate() {
            mem::replace(&mut *r.receiver.borrow_mut(), mem::replace(&mut receivers[index], None));
        }

        if r_id == -1 {
            return Err(MpscError::UnknownError);
        }

        let receivers = &mut self.receivers;
        match receivers[r_index].recv() {
            Ok((data, channels, shmems)) =>
                Ok(vec![OsIpcSelectionResult::DataReceived(r_id, data, channels, shmems)]),
            Err(MpscError::ChannelClosedError) => {
                receivers.remove(r_index);
                self.receiver_ids.remove(r_index);
                Ok(vec![OsIpcSelectionResult::ChannelClosed(r_id)])
            },
            Err(err) => Err(err),
        }
    }
}

pub enum OsIpcSelectionResult {
    DataReceived(i64, Vec<u8>, Vec<OsOpaqueIpcChannel>, Vec<OsIpcSharedMemory>),
    ChannelClosed(i64),
}

impl OsIpcSelectionResult {
    pub fn unwrap(self) -> (i64, Vec<u8>, Vec<OsOpaqueIpcChannel>, Vec<OsIpcSharedMemory>) {
        match self {
            OsIpcSelectionResult::DataReceived(id, data, channels, shared_memory_regions) => {
                (id, data, channels, shared_memory_regions)
            }
            OsIpcSelectionResult::ChannelClosed(id) => {
                panic!("OsIpcSelectionResult::unwrap(): receiver ID {} was closed!", id)
            }
        }
    }
}

pub struct OsIpcOneShotServer {
    receiver: RefCell<Option<OsIpcReceiver>>,
    name: String,
}

impl OsIpcOneShotServer {
    pub fn new() -> Result<(OsIpcOneShotServer, String),MpscError> {
        let (sender, receiver) = match channel() {
            Ok((s,r)) => (s,r),
            Err(err) => return Err(err),
        };

        let name = Uuid::new_v4().to_string();
        let record = ServerRecord::new(sender);
        ONE_SHOT_SERVERS.lock().unwrap().insert(name.clone(), record);
        Ok((OsIpcOneShotServer {
            receiver: RefCell::new(Some(receiver)),
            name: name.clone(),
        },name.clone()))
    }

    pub fn accept(&self) -> Result<(OsIpcReceiver,
                                    Vec<u8>,
                                    Vec<OsOpaqueIpcChannel>,
                                    Vec<OsIpcSharedMemory>),MpscError>
    {
        ONE_SHOT_SERVERS.lock().unwrap().get(&self.name).unwrap().accept();
        let receiver = self.receiver.borrow_mut().take().unwrap();
        let (data, channels, shmems) = receiver.recv().unwrap();
        Ok((receiver, data, channels, shmems))
    }
}

#[derive(PartialEq, Debug)]
pub enum OsIpcChannel {
    Sender(OsIpcSender),
    Receiver(OsIpcReceiver),
}

#[derive(PartialEq, Debug)]
pub struct OsOpaqueIpcChannel {
    channel: RefCell<Option<OsIpcChannel>>,
}

impl OsOpaqueIpcChannel {
    fn new(channel: OsIpcChannel) -> OsOpaqueIpcChannel {
        OsOpaqueIpcChannel {
            channel: RefCell::new(Some(channel))
        }
    }

    pub fn to_receiver(&self) -> OsIpcReceiver {
        match self.channel.borrow_mut().take().unwrap() {
            OsIpcChannel::Sender(_) => panic!("Opaque channel is not a receiver!"),
            OsIpcChannel::Receiver(r) => r
        }
    }

    pub fn to_sender(&self) -> OsIpcSender {
        match self.channel.borrow_mut().take().unwrap() {
            OsIpcChannel::Sender(s) => s,
            OsIpcChannel::Receiver(_) => panic!("Opaque channel is not a sender!"),
        }
    }
}

pub struct OsIpcSharedMemory {
    ptr: *mut u8,
    length: usize,
    data: Arc<Vec<u8>>,
}

unsafe impl Send for OsIpcSharedMemory {}
unsafe impl Sync for OsIpcSharedMemory {}

impl Clone for OsIpcSharedMemory {
    fn clone(&self) -> OsIpcSharedMemory {
        OsIpcSharedMemory {
            ptr: self.ptr,
            length: self.length,
            data: self.data.clone(),
        }
    }
}

impl PartialEq for OsIpcSharedMemory {
    fn eq(&self, other: &OsIpcSharedMemory) -> bool {
        **self == **other
    }
}

impl Debug for OsIpcSharedMemory {
    fn fmt(&self, formatter: &mut Formatter) -> Result<(), fmt::Error> {
        (**self).fmt(formatter)
    }
}

impl Deref for OsIpcSharedMemory {
    type Target = [u8];

    #[inline]
    fn deref(&self) -> &[u8] {
        if self.ptr.is_null() {
            panic!("attempted to access a consumed `OsIpcSharedMemory`")
        }
        unsafe {
            slice::from_raw_parts(self.ptr, self.length)
        }
    }
}

impl OsIpcSharedMemory {
    pub fn from_byte(byte: u8, length: usize) -> OsIpcSharedMemory {
        let mut v = Arc::new(vec![byte; length]);
        OsIpcSharedMemory {
            ptr: Arc::get_mut(&mut v).unwrap().as_mut_ptr(),
            length: length,
            data: v
        }
    }

    pub fn from_bytes(bytes: &[u8]) -> OsIpcSharedMemory {
        let mut v = Arc::new(bytes.to_vec());
        OsIpcSharedMemory {
            ptr: Arc::get_mut(&mut v).unwrap().as_mut_ptr(),
            length: v.len(),
            data: v
        }
    }
}

#[derive(Debug, PartialEq)]
pub enum MpscError {
    ChannelClosedError,
    UnknownError,
}

impl MpscError {
    #[allow(dead_code)]
    pub fn channel_is_closed(&self) -> bool {
        *self == MpscError::ChannelClosedError
    }
}

impl From<MpscError> for DeserializeError {
    fn from(mpsc_error: MpscError) -> DeserializeError {
        DeserializeError::IoError(mpsc_error.into())
    }
}

impl From<MpscError> for Error {
    fn from(mpsc_error: MpscError) -> Error {
        match mpsc_error {
            MpscError::ChannelClosedError => {
                Error::new(ErrorKind::BrokenPipe, "MPSC channel closed")
            }
            MpscError::UnknownError => Error::new(ErrorKind::Other, "Other MPSC channel error"),
        }
    }
}

<commit_msg>Auto merge of #89 - servo:threadsafe, r=Wafflespeanut<commit_after>// Copyright 2015 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use bincode::serde::DeserializeError;
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
use std::collections::hash_map::HashMap;
use std::cell::{RefCell};
use std::io::{Error, ErrorKind};
use std::slice;
use std::fmt::{self, Debug, Formatter};
use std::cmp::{PartialEq};
use std::ops::Deref;
use std::mem;
use std::usize;

use uuid::Uuid;

struct ServerRecord {
    sender: OsIpcSender,
    conn_sender: mpsc::Sender<bool>,
    conn_receiver: Mutex<mpsc::Receiver<bool>>,
}

impl ServerRecord {
    fn new(sender: OsIpcSender) -> ServerRecord {
        let (tx, rx) = mpsc::channel::<bool>();
        ServerRecord {
            sender: sender,
            conn_sender: tx,
            conn_receiver: Mutex::new(rx),
        }
    }

    fn accept(&self) {
        self.conn_receiver.lock().unwrap().recv().unwrap();
    }

    fn connect(&self) {
        self.conn_sender.send(true).unwrap();
    }
}

lazy_static! {
    static ref ONE_SHOT_SERVERS: Mutex<HashMap<String,ServerRecord>> = Mutex::new(HashMap::new());
}

struct MpscChannelMessage(Vec<u8>, Vec<OsIpcChannel>, Vec<OsIpcSharedMemory>);

pub fn channel() -> Result<(OsIpcSender, OsIpcReceiver),MpscError> {
    let (base_sender, base_receiver) = mpsc::channel::<MpscChannelMessage>();
    Ok((OsIpcSender::new(base_sender), OsIpcReceiver::new(base_receiver)))
}

pub struct OsIpcReceiver {
    receiver: Arc<Mutex<Option<mpsc::Receiver<MpscChannelMessage>>>>,
}

impl PartialEq for OsIpcReceiver {
    fn eq(&self, other: &OsIpcReceiver) -> bool {
        self.receiver.lock().unwrap().as_ref().map(|rx| rx as *const _) ==
            other.receiver.lock().unwrap().as_ref().map(|rx| rx as *const _)
    }
}

// Can't derive, as mpsc::Receiver doesn't implement Debug.
impl fmt::Debug for OsIpcReceiver {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // Not sure there is anything useful we could print here.
        write!(f, "OsIpcReceiver {{ .. }}")
    }
}

impl OsIpcReceiver {
    fn new(receiver: mpsc::Receiver<MpscChannelMessage>) -> OsIpcReceiver {
        OsIpcReceiver {
            receiver: Arc::new(Mutex::new(Some(receiver))),
        }
    }

    pub fn consume(&self) -> OsIpcReceiver {
        let receiver = self.receiver.lock().unwrap().take();
        OsIpcReceiver::new(receiver.unwrap())
    }

    pub fn recv(&self) -> Result<(Vec<u8>, Vec<OsOpaqueIpcChannel>, Vec<OsIpcSharedMemory>),MpscError> {
        let r = self.receiver.lock().unwrap();
        match r.as_ref().unwrap().recv() {
            Ok(MpscChannelMessage(d,c,s)) => Ok((d,
                                                 c.into_iter().map(OsOpaqueIpcChannel::new).collect(),
                                                 s)),
            Err(_) => Err(MpscError::ChannelClosedError),
        }
    }

    pub fn try_recv(&self) -> Result<(Vec<u8>, Vec<OsOpaqueIpcChannel>, Vec<OsIpcSharedMemory>),MpscError> {
        let r = self.receiver.lock().unwrap();
        match r.as_ref().unwrap().try_recv() {
            Ok(MpscChannelMessage(d,c,s)) => Ok((d,
                                                 c.into_iter().map(OsOpaqueIpcChannel::new).collect(),
                                                 s)),
            Err(_) => Err(MpscError::ChannelClosedError),
        }
    }
}

#[derive(Clone)]
pub struct OsIpcSender {
    sender: Arc<Mutex<mpsc::Sender<MpscChannelMessage>>>,
}

impl PartialEq for OsIpcSender {
    fn eq(&self, other: &OsIpcSender) -> bool {
        &*self.sender.lock().unwrap() as *const _ ==
            &*other.sender.lock().unwrap() as *const _
    }
}

// Can't derive, as mpsc::Sender doesn't implement Debug.
impl fmt::Debug for OsIpcSender {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // Not sure there is anything useful we could print here.
        write!(f, "OsIpcSender {{ .. }}")
    }
}

impl OsIpcSender {
    fn new(sender: mpsc::Sender<MpscChannelMessage>) -> OsIpcSender {
        OsIpcSender {
            sender: Arc::new(Mutex::new(sender)),
        }
    }

    pub fn connect(name: String) -> Result<OsIpcSender,MpscError> {
        let record = ONE_SHOT_SERVERS.lock().unwrap().remove(&name).unwrap();
        record.connect();
        Ok(record.sender)
    }

    pub fn get_max_fragment_size() -> usize {
        usize::MAX
    }

    pub fn send(&self,
                data: &[u8],
                ports: Vec<OsIpcChannel>,
                shared_memory_regions: Vec<OsIpcSharedMemory>)
                -> Result<(),MpscError>
    {
        match self.sender.lock().unwrap().send(MpscChannelMessage(data.to_vec(), ports, shared_memory_regions)) {
            Err(_) => Err(MpscError::ChannelClosedError),
            Ok(_) => Ok(()),
        }
    }
}

pub struct OsIpcReceiverSet {
    last_index: usize,
    receiver_ids: Vec<usize>,
    receivers: Vec<OsIpcReceiver>,
}

impl OsIpcReceiverSet {
    pub fn new() -> Result<OsIpcReceiverSet,MpscError> {
        Ok(OsIpcReceiverSet {
            last_index: 0,
            receiver_ids: vec![],
            receivers: vec![],
        })
    }

    pub fn add(&mut self, receiver: OsIpcReceiver) -> Result<i64,MpscError> {
        self.last_index += 1;
        self.receiver_ids.push(self.last_index);
        self.receivers.push(receiver.consume());
        Ok(self.last_index as i64)
    }

    pub fn select(&mut self) -> Result<Vec<OsIpcSelectionResult>,MpscError> {
        let mut receivers: Vec<Option<mpsc::Receiver<MpscChannelMessage>>> = Vec::with_capacity(self.receivers.len());
        let mut r_id: i64 = -1;
        let mut r_index: usize = 0;

        {
            let select = mpsc::Select::new();
            // we *must* allocate exact capacity for this, because the Handles *can't move*
            let mut handles: Vec<mpsc::Handle<MpscChannelMessage>> = Vec::with_capacity(self.receivers.len());

            for r in &self.receivers {
                let inner_r = mem::replace(&mut *r.receiver.lock().unwrap(), None);
                receivers.push(inner_r);
            }

            for r in &receivers {
                unsafe {
                    handles.push(select.handle(r.as_ref().unwrap()));
                    handles.last_mut().unwrap().add();
                }
            }

            let id = select.wait();

            for (index,h) in handles.iter().enumerate() {
                if h.id() == id {
                    r_index = index;
                    r_id = self.receiver_ids[index] as i64;
                    break;
                }
            }
        }

        // put the receivers back
        for (index,r) in self.receivers.iter().enumerate() {
            mem::replace(&mut *r.receiver.lock().unwrap(), mem::replace(&mut receivers[index], None));
        }

        if r_id == -1 {
            return Err(MpscError::UnknownError);
        }

        let receivers = &mut self.receivers;
        match receivers[r_index].recv() {
            Ok((data, channels, shmems)) =>
                Ok(vec![OsIpcSelectionResult::DataReceived(r_id, data, channels, shmems)]),
            Err(MpscError::ChannelClosedError) => {
                receivers.remove(r_index);
                self.receiver_ids.remove(r_index);
                Ok(vec![OsIpcSelectionResult::ChannelClosed(r_id)])
            },
            Err(err) => Err(err),
        }
    }
}

pub enum OsIpcSelectionResult {
    DataReceived(i64, Vec<u8>, Vec<OsOpaqueIpcChannel>, Vec<OsIpcSharedMemory>),
    ChannelClosed(i64),
}

impl OsIpcSelectionResult {
    pub fn unwrap(self) -> (i64, Vec<u8>, Vec<OsOpaqueIpcChannel>, Vec<OsIpcSharedMemory>) {
        match self {
            OsIpcSelectionResult::DataReceived(id, data, channels, shared_memory_regions) => {
                (id, data, channels, shared_memory_regions)
            }
            OsIpcSelectionResult::ChannelClosed(id) => {
                panic!("OsIpcSelectionResult::unwrap(): receiver ID {} was closed!", id)
            }
        }
    }
}

pub struct OsIpcOneShotServer {
    receiver: RefCell<Option<OsIpcReceiver>>,
    name: String,
}

impl OsIpcOneShotServer {
    pub fn new() -> Result<(OsIpcOneShotServer, String),MpscError> {
        let (sender, receiver) = match channel() {
            Ok((s,r)) => (s,r),
            Err(err) => return Err(err),
        };

        let name = Uuid::new_v4().to_string();
        let record = ServerRecord::new(sender);
        ONE_SHOT_SERVERS.lock().unwrap().insert(name.clone(), record);
        Ok((OsIpcOneShotServer {
            receiver: RefCell::new(Some(receiver)),
            name: name.clone(),
        },name.clone()))
    }

    pub fn accept(&self) -> Result<(OsIpcReceiver,
                                    Vec<u8>,
                                    Vec<OsOpaqueIpcChannel>,
                                    Vec<OsIpcSharedMemory>),MpscError>
    {
        ONE_SHOT_SERVERS.lock().unwrap().get(&self.name).unwrap().accept();
        let receiver = self.receiver.borrow_mut().take().unwrap();
        let (data, channels, shmems) = receiver.recv().unwrap();
        Ok((receiver, data, channels, shmems))
    }
}

#[derive(PartialEq, Debug)]
pub enum OsIpcChannel {
    Sender(OsIpcSender),
    Receiver(OsIpcReceiver),
}

#[derive(PartialEq, Debug)]
pub struct OsOpaqueIpcChannel {
    channel: RefCell<Option<OsIpcChannel>>,
}

impl OsOpaqueIpcChannel {
    fn new(channel: OsIpcChannel) -> OsOpaqueIpcChannel {
        OsOpaqueIpcChannel {
            channel: RefCell::new(Some(channel))
        }
    }

    pub fn to_receiver(&self) -> OsIpcReceiver {
        match self.channel.borrow_mut().take().unwrap() {
            OsIpcChannel::Sender(_) => panic!("Opaque channel is not a receiver!"),
            OsIpcChannel::Receiver(r) => r
        }
    }

    pub fn to_sender(&self) -> OsIpcSender {
        match self.channel.borrow_mut().take().unwrap() {
            OsIpcChannel::Sender(s) => s,
            OsIpcChannel::Receiver(_) => panic!("Opaque channel is not a sender!"),
        }
    }
}

pub struct OsIpcSharedMemory {
    ptr: *mut u8,
    length: usize,
    data: Arc<Vec<u8>>,
}

unsafe impl Send for OsIpcSharedMemory {}
unsafe impl Sync for OsIpcSharedMemory {}

impl Clone for OsIpcSharedMemory {
    fn clone(&self) -> OsIpcSharedMemory {
        OsIpcSharedMemory {
            ptr: self.ptr,
            length: self.length,
            data: self.data.clone(),
        }
    }
}

impl PartialEq for OsIpcSharedMemory {
    fn eq(&self, other: &OsIpcSharedMemory) -> bool {
        **self == **other
    }
}

impl Debug for OsIpcSharedMemory {
    fn fmt(&self, formatter: &mut Formatter) -> Result<(), fmt::Error> {
        (**self).fmt(formatter)
    }
}

impl Deref for OsIpcSharedMemory {
    type Target = [u8];

    #[inline]
    fn deref(&self) -> &[u8] {
        if self.ptr.is_null() {
            panic!("attempted to access a consumed `OsIpcSharedMemory`")
        }
        unsafe {
            slice::from_raw_parts(self.ptr, self.length)
        }
    }
}

impl OsIpcSharedMemory {
    pub fn from_byte(byte: u8, length: usize) -> OsIpcSharedMemory {
        let mut v = Arc::new(vec![byte; length]);
        OsIpcSharedMemory {
            ptr: Arc::get_mut(&mut v).unwrap().as_mut_ptr(),
            length: length,
            data: v
        }
    }

    pub fn from_bytes(bytes: &[u8]) -> OsIpcSharedMemory {
        let mut v = Arc::new(bytes.to_vec());
        OsIpcSharedMemory {
            ptr: Arc::get_mut(&mut v).unwrap().as_mut_ptr(),
            length: v.len(),
            data: v
        }
    }
}

#[derive(Debug, PartialEq)]
pub enum MpscError {
    ChannelClosedError,
    UnknownError,
}

impl MpscError {
    #[allow(dead_code)]
    pub fn channel_is_closed(&self) -> bool {
        *self == MpscError::ChannelClosedError
    }
}

impl From<MpscError> for DeserializeError {
    fn from(mpsc_error: MpscError) -> DeserializeError {
        DeserializeError::IoError(mpsc_error.into())
    }
}

impl From<MpscError> for Error {
    fn from(mpsc_error: MpscError) -> Error {
        match mpsc_error {
            MpscError::ChannelClosedError => {
                Error::new(ErrorKind::BrokenPipe, "MPSC channel closed")
            }
            MpscError::UnknownError => Error::new(ErrorKind::Other, "Other MPSC channel error"),
        }
    }
}

<|endoftext|>


<commit_before>use std::cmp::Ordering;
use std::marker::ContravariantLifetime;
use std::ops::Index;

use objc::runtime::Object;
use objc::{Id, IdSlice, IntoIdVector, Owned, Ownership, Shared, ShareId};

use {INSCopying, INSMutableCopying, INSObject};

#[repr(isize)]
#[derive(Copy)]
pub enum NSComparisonResult {
    Ascending  = -1,
    Same       = 0,
    Descending = 1,
}

impl NSComparisonResult {
    pub fn from_ordering(order: Ordering) -> NSComparisonResult {
        match order {
            Ordering::Less => NSComparisonResult::Ascending,
            Ordering::Equal => NSComparisonResult::Same,
            Ordering::Greater => NSComparisonResult::Descending,
        }
    }

    pub fn as_ordering(&self) -> Ordering {
        match *self {
            NSComparisonResult::Ascending => Ordering::Less,
            NSComparisonResult::Same => Ordering::Equal,
            NSComparisonResult::Descending => Ordering::Greater,
        }
    }
}

#[repr(C)]
#[derive(Copy)]
pub struct NSRange {
    pub location: usize,
    pub length: usize,
}

pub struct NSEnumerator<'a, T> {
    id: Id<Object>,
    marker: ContravariantLifetime<'a>,
}

impl<'a, T: INSObject> NSEnumerator<'a, T> {
    pub unsafe fn from_ptr(ptr: *mut Object) -> NSEnumerator<'a, T> {
        NSEnumerator { id: Id::from_ptr(ptr), marker: ContravariantLifetime }
    }
}

impl<'a, T: INSObject> Iterator for NSEnumerator<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<&'a T> {
        unsafe {
            let obj: *mut T = msg_send![self.id, nextObject];
            obj.as_ref()
        }
    }
}

pub trait INSArray : INSObject {
    type Item: INSObject;
    type Own: Ownership;

    fn count(&self) -> usize {
        unsafe {
            msg_send![self, count]
        }
    }

    fn object_at(&self, index: usize) -> &Self::Item {
        unsafe {
            let obj: *const Self::Item = msg_send![self, objectAtIndex:index];
            &*obj
        }
    }

    fn first_object(&self) -> Option<&Self::Item> {
        unsafe {
            let obj: *const Self::Item = msg_send![self, firstObject];
            obj.as_ref()
        }
    }

    fn last_object(&self) -> Option<&Self::Item> {
        unsafe {
            let obj: *const Self::Item = msg_send![self, lastObject];
            obj.as_ref()
        }
    }

    fn object_enumerator(&self) -> NSEnumerator<Self::Item> {
        unsafe {
            let result: *mut Object = msg_send![self, objectEnumerator];
            NSEnumerator::from_ptr(result)
        }
    }

    unsafe fn from_refs(refs: &[&Self::Item]) -> Id<Self> {
        let cls = <Self as INSObject>::class();
        let obj: *mut Self = msg_send![cls, alloc];
        let obj: *mut Self = msg_send![obj, initWithObjects:refs.as_ptr()
                                                      count:refs.len()];
        Id::from_retained_ptr(obj)
    }

    fn from_vec(vec: Vec<Id<Self::Item, Self::Own>>) -> Id<Self> {
        let refs = vec.as_refs_slice();
        unsafe {
            INSArray::from_refs(refs)
        }
    }

    fn objects_in_range(&self, start: usize, len: usize) -> Vec<&Self::Item> {
        let mut vec: Vec<&Self::Item> = Vec::with_capacity(len);
        let range = NSRange { location: start, length: len };
        unsafe {
            let _: () = msg_send![self, getObjects:vec.as_ptr() range:range];
            vec.set_len(len);
        }
        vec
    }

    fn to_vec(&self) -> Vec<&Self::Item> {
        self.objects_in_range(0, self.count())
    }

    fn into_vec(array: Id<Self>) -> Vec<Id<Self::Item, Self::Own>> {
        let vec = array.to_vec();
        unsafe {
            vec.into_id_vec()
        }
    }
}

pub trait INSOwnedArray : INSArray<Own=Owned> {
    fn mut_object_at(&mut self, index: usize) -> &mut Self::Item {
        unsafe {
            let result: *mut Self::Item = msg_send![self, objectAtIndex:index];
            &mut *result
        }
    }
}

pub trait INSSharedArray : INSArray<Own=Shared> {
    fn shared_object_at(&self, index: usize) -> ShareId<Self::Item> {
        let obj = self.object_at(index);
        unsafe {
            Id::from_ptr(obj as *const _ as *mut Self::Item)
        }
    }

    fn from_slice(slice: &[ShareId<Self::Item>]) -> Id<Self> {
        let refs = slice.as_refs_slice();
        unsafe {
            INSArray::from_refs(refs)
        }
    }

    fn to_shared_vec(&self) -> Vec<ShareId<Self::Item>> {
        let vec = self.to_vec();
        unsafe {
            vec.into_id_vec()
        }
    }
}

#[allow(missing_copy_implementations)]
pub enum NSArray<T, O = Owned> { }

object_impl!(NSArray<T, O>);

impl<T: INSObject, O: Ownership> INSArray for NSArray<T, O> {
    type Item = T;
    type Own = O;
}

impl<T: INSObject> INSOwnedArray for NSArray<T, Owned> { }

impl<T: INSObject> INSSharedArray for NSArray<T, Shared> { }

impl<T> INSCopying for NSArray<T, Shared> {
    type Output = NSSharedArray<T>;
}

impl<T> INSMutableCopying for NSArray<T, Shared> {
    type Output = NSMutableSharedArray<T>;
}

impl<T: INSObject, O: Ownership> Index<usize> for NSArray<T, O> {
    type Output = T;

    fn index(&self, index: &usize) -> &T {
        self.object_at(*index)
    }
}

pub type NSSharedArray<T> = NSArray<T, Shared>;

pub trait INSMutableArray : INSArray {
    fn add_object(&mut self, obj: Id<Self::Item, Self::Own>) {
        unsafe {
            let _: () = msg_send![self, addObject:obj];
        }
    }

    fn insert_object_at(&mut self, index: usize, obj: Id<Self::Item, Self::Own>) {
        unsafe {
            let _: () = msg_send![self, insertObject:obj atIndex:index];
        }
    }

    fn replace_object_at(&mut self, index: usize, obj: Id<Self::Item, Self::Own>) ->
            Id<Self::Item, Self::Own> {
        let old_obj = unsafe {
            let obj = self.object_at(index);
            Id::from_ptr(obj as *const _ as *mut Self::Item)
        };
        unsafe {
            let _: () = msg_send![self, replaceObjectAtIndex:index
                                                  withObject:obj];
        }
        old_obj
    }

    fn remove_object_at(&mut self, index: usize) -> Id<Self::Item, Self::Own> {
        let obj = unsafe {
            let obj = self.object_at(index);
            Id::from_ptr(obj as *const _ as *mut Self::Item)
        };
        unsafe {
            let _: () = msg_send![self, removeObjectAtIndex:index];
        }
        obj
    }

    fn remove_last_object(&mut self) -> Id<Self::Item, Self::Own> {
        let obj = self.last_object().map(|obj| unsafe {
            Id::from_ptr(obj as *const _ as *mut Self::Item)
        });
        unsafe {
            let _: () = msg_send![self, removeLastObject];
        }
        // removeLastObject would have failed if the array is empty,
        // so we know this won't be None
        obj.unwrap()
    }

    fn remove_all_objects(&mut self) {
        unsafe {
            let _: () = msg_send![self, removeAllObjects];
        }
    }

    fn sort_by<F: FnMut(&Self::Item, &Self::Item) -> Ordering>(&mut self, compare: F) {
        extern fn compare_with_closure<T, F: FnMut(&T, &T) -> Ordering>(
                obj1: &T, obj2: &T, compare: &mut F) -> NSComparisonResult {
            NSComparisonResult::from_ordering((*compare)(obj1, obj2))
        }

        let mut closure = compare;
        unsafe {
            let _: () = msg_send![self, sortUsingFunction:compare_with_closure::<Self::Item, F>
                                                  context:&mut closure];
        }
    }
}

#[allow(missing_copy_implementations)]
pub enum NSMutableArray<T, O = Owned> { }

object_impl!(NSMutableArray<T, O>);

impl<T: INSObject, O: Ownership> INSArray for NSMutableArray<T, O> {
    type Item = T;
    type Own = O;
}

impl<T: INSObject> INSOwnedArray for NSMutableArray<T, Owned> { }

impl<T: INSObject> INSSharedArray for NSMutableArray<T, Shared> { }

impl<T: INSObject, O: Ownership> INSMutableArray for NSMutableArray<T, O> { }

impl<T> INSCopying for NSMutableArray<T, Shared> {
    type Output = NSSharedArray<T>;
}

impl<T> INSMutableCopying for NSMutableArray<T, Shared> {
    type Output = NSMutableSharedArray<T>;
}

impl<T: INSObject, O: Ownership> Index<usize> for NSMutableArray<T, O> {
    type Output = T;

    fn index(&self, index: &usize) -> &T {
        self.object_at(*index)
    }
}

pub type NSMutableSharedArray<T> = NSMutableArray<T, Shared>;

#[cfg(test)]
mod tests {
    use objc::{Id};
    use {INSObject, INSString, NSObject, NSString};
    use super::{INSArray, INSMutableArray, NSArray, NSMutableArray};

    fn sample_array(len: usize) -> Id<NSArray<NSObject>> {
        let mut vec: Vec<Id<NSObject>> = Vec::with_capacity(len);
        for _ in 0..len {
            vec.push(INSObject::new());
        }
        INSArray::from_vec(vec)
    }

    #[test]
    fn test_count() {
        let empty_array: Id<NSArray<NSObject>> = INSObject::new();
        assert!(empty_array.count() == 0);

        let array = sample_array(4);
        assert!(array.count() == 4);
    }

    #[test]
    fn test_object_at() {
        let array = sample_array(4);
        assert!(array.object_at(0) != array.object_at(3));
        assert!(array.first_object().unwrap() == array.object_at(0));
        assert!(array.last_object().unwrap() == array.object_at(3));

        let empty_array: Id<NSArray<NSObject>> = INSObject::new();
        assert!(empty_array.first_object().is_none());
        assert!(empty_array.last_object().is_none());
    }

    #[test]
    fn test_object_enumerator() {
        let array = sample_array(4);

        assert!(array.object_enumerator().count() == 4);
        assert!(array.object_enumerator()
                     .enumerate()
                     .all(|(i, obj)| obj == array.object_at(i)));
    }

    #[test]
    fn test_objects_in_range() {
        let array = sample_array(4);

        let middle_objs = array.objects_in_range(1, 2);
        assert!(middle_objs.len() == 2);
        assert!(middle_objs[0] == array.object_at(1));
        assert!(middle_objs[1] == array.object_at(2));

        let empty_objs = array.objects_in_range(1, 0);
        assert!(empty_objs.len() == 0);

        let all_objs = array.objects_in_range(0, 4);
        assert!(all_objs.len() == 4);
    }

    #[test]
    fn test_into_vec() {
        let array = sample_array(4);

        let vec = INSArray::into_vec(array);
        assert!(vec.len() == 4);
    }

    #[test]
    fn test_add_object() {
        let mut array: Id<NSMutableArray<NSObject>> = INSObject::new();
        let obj: Id<NSObject> = INSObject::new();
        array.add_object(obj);

        assert!(array.count() == 1);
        assert!(array.object_at(0) == array.object_at(0));

        let obj: Id<NSObject> = INSObject::new();
        array.insert_object_at(0, obj);
        assert!(array.count() == 2);
    }

    #[test]
    fn test_replace_object() {
        let mut array: Id<NSMutableArray<NSObject>> = INSObject::new();
        let obj: Id<NSObject> = INSObject::new();
        array.add_object(obj);

        let obj: Id<NSObject> = INSObject::new();
        let old_obj = array.replace_object_at(0, obj);
        assert!(&*old_obj != array.object_at(0));
    }

    #[test]
    fn test_remove_object() {
        let mut array: Id<NSMutableArray<NSObject>> = INSObject::new();
        for _ in 0..4 {
            let obj: Id<NSObject> = INSObject::new();
            array.add_object(obj);
        }

        array.remove_object_at(1);
        assert!(array.count() == 3);

        array.remove_last_object();
        assert!(array.count() == 2);

        array.remove_all_objects();
        assert!(array.count() == 0);
    }

    #[test]
    fn test_sort() {
        let strings: Vec<Id<NSString>> = vec![
            INSString::from_str("hello"),
            INSString::from_str("hi"),
        ];
        let mut strings: Id<NSMutableArray<_>> = INSArray::from_vec(strings);

        strings.sort_by(|s1, s2| s1.as_str().len().cmp(&s2.as_str().len()));
        assert!(strings[0].as_str() == "hi");
        assert!(strings[1].as_str() == "hello");
    }
}
<commit_msg>Use range notation for NSArray's objects_in_range.<commit_after>use std::cmp::Ordering;
use std::marker::ContravariantLifetime;
use std::ops::{Index, Range};

use objc::runtime::Object;
use objc::{Id, IdSlice, IntoIdVector, Owned, Ownership, Shared, ShareId};

use {INSCopying, INSMutableCopying, INSObject};

#[repr(isize)]
#[derive(Copy)]
pub enum NSComparisonResult {
    Ascending  = -1,
    Same       = 0,
    Descending = 1,
}

impl NSComparisonResult {
    pub fn from_ordering(order: Ordering) -> NSComparisonResult {
        match order {
            Ordering::Less => NSComparisonResult::Ascending,
            Ordering::Equal => NSComparisonResult::Same,
            Ordering::Greater => NSComparisonResult::Descending,
        }
    }

    pub fn as_ordering(&self) -> Ordering {
        match *self {
            NSComparisonResult::Ascending => Ordering::Less,
            NSComparisonResult::Same => Ordering::Equal,
            NSComparisonResult::Descending => Ordering::Greater,
        }
    }
}

#[repr(C)]
#[derive(Copy)]
pub struct NSRange {
    pub location: usize,
    pub length: usize,
}

impl NSRange {
    pub fn from_range(range: Range<usize>) -> NSRange {
        assert!(range.end >= range.start);
        NSRange { location: range.start, length: range.end - range.start }
    }

    pub fn as_range(&self) -> Range<usize> {
        Range { start: self.location, end: self.location + self.length }
    }
}

pub struct NSEnumerator<'a, T> {
    id: Id<Object>,
    marker: ContravariantLifetime<'a>,
}

impl<'a, T: INSObject> NSEnumerator<'a, T> {
    pub unsafe fn from_ptr(ptr: *mut Object) -> NSEnumerator<'a, T> {
        NSEnumerator { id: Id::from_ptr(ptr), marker: ContravariantLifetime }
    }
}

impl<'a, T: INSObject> Iterator for NSEnumerator<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<&'a T> {
        unsafe {
            let obj: *mut T = msg_send![self.id, nextObject];
            obj.as_ref()
        }
    }
}

pub trait INSArray : INSObject {
    type Item: INSObject;
    type Own: Ownership;

    fn count(&self) -> usize {
        unsafe {
            msg_send![self, count]
        }
    }

    fn object_at(&self, index: usize) -> &Self::Item {
        unsafe {
            let obj: *const Self::Item = msg_send![self, objectAtIndex:index];
            &*obj
        }
    }

    fn first_object(&self) -> Option<&Self::Item> {
        unsafe {
            let obj: *const Self::Item = msg_send![self, firstObject];
            obj.as_ref()
        }
    }

    fn last_object(&self) -> Option<&Self::Item> {
        unsafe {
            let obj: *const Self::Item = msg_send![self, lastObject];
            obj.as_ref()
        }
    }

    fn object_enumerator(&self) -> NSEnumerator<Self::Item> {
        unsafe {
            let result: *mut Object = msg_send![self, objectEnumerator];
            NSEnumerator::from_ptr(result)
        }
    }

    unsafe fn from_refs(refs: &[&Self::Item]) -> Id<Self> {
        let cls = <Self as INSObject>::class();
        let obj: *mut Self = msg_send![cls, alloc];
        let obj: *mut Self = msg_send![obj, initWithObjects:refs.as_ptr()
                                                      count:refs.len()];
        Id::from_retained_ptr(obj)
    }

    fn from_vec(vec: Vec<Id<Self::Item, Self::Own>>) -> Id<Self> {
        let refs = vec.as_refs_slice();
        unsafe {
            INSArray::from_refs(refs)
        }
    }

    fn objects_in_range(&self, range: Range<usize>) -> Vec<&Self::Item> {
        let range = NSRange::from_range(range);
        let mut vec: Vec<&Self::Item> = Vec::with_capacity(range.length);
        unsafe {
            let _: () = msg_send![self, getObjects:vec.as_ptr() range:range];
            vec.set_len(range.length);
        }
        vec
    }

    fn to_vec(&self) -> Vec<&Self::Item> {
        self.objects_in_range(0..self.count())
    }

    fn into_vec(array: Id<Self>) -> Vec<Id<Self::Item, Self::Own>> {
        let vec = array.to_vec();
        unsafe {
            vec.into_id_vec()
        }
    }
}

pub trait INSOwnedArray : INSArray<Own=Owned> {
    fn mut_object_at(&mut self, index: usize) -> &mut Self::Item {
        unsafe {
            let result: *mut Self::Item = msg_send![self, objectAtIndex:index];
            &mut *result
        }
    }
}

pub trait INSSharedArray : INSArray<Own=Shared> {
    fn shared_object_at(&self, index: usize) -> ShareId<Self::Item> {
        let obj = self.object_at(index);
        unsafe {
            Id::from_ptr(obj as *const _ as *mut Self::Item)
        }
    }

    fn from_slice(slice: &[ShareId<Self::Item>]) -> Id<Self> {
        let refs = slice.as_refs_slice();
        unsafe {
            INSArray::from_refs(refs)
        }
    }

    fn to_shared_vec(&self) -> Vec<ShareId<Self::Item>> {
        let vec = self.to_vec();
        unsafe {
            vec.into_id_vec()
        }
    }
}

#[allow(missing_copy_implementations)]
pub enum NSArray<T, O = Owned> { }

object_impl!(NSArray<T, O>);

impl<T: INSObject, O: Ownership> INSArray for NSArray<T, O> {
    type Item = T;
    type Own = O;
}

impl<T: INSObject> INSOwnedArray for NSArray<T, Owned> { }

impl<T: INSObject> INSSharedArray for NSArray<T, Shared> { }

impl<T> INSCopying for NSArray<T, Shared> {
    type Output = NSSharedArray<T>;
}

impl<T> INSMutableCopying for NSArray<T, Shared> {
    type Output = NSMutableSharedArray<T>;
}

impl<T: INSObject, O: Ownership> Index<usize> for NSArray<T, O> {
    type Output = T;

    fn index(&self, index: &usize) -> &T {
        self.object_at(*index)
    }
}

pub type NSSharedArray<T> = NSArray<T, Shared>;

pub trait INSMutableArray : INSArray {
    fn add_object(&mut self, obj: Id<Self::Item, Self::Own>) {
        unsafe {
            let _: () = msg_send![self, addObject:obj];
        }
    }

    fn insert_object_at(&mut self, index: usize, obj: Id<Self::Item, Self::Own>) {
        unsafe {
            let _: () = msg_send![self, insertObject:obj atIndex:index];
        }
    }

    fn replace_object_at(&mut self, index: usize, obj: Id<Self::Item, Self::Own>) ->
            Id<Self::Item, Self::Own> {
        let old_obj = unsafe {
            let obj = self.object_at(index);
            Id::from_ptr(obj as *const _ as *mut Self::Item)
        };
        unsafe {
            let _: () = msg_send![self, replaceObjectAtIndex:index
                                                  withObject:obj];
        }
        old_obj
    }

    fn remove_object_at(&mut self, index: usize) -> Id<Self::Item, Self::Own> {
        let obj = unsafe {
            let obj = self.object_at(index);
            Id::from_ptr(obj as *const _ as *mut Self::Item)
        };
        unsafe {
            let _: () = msg_send![self, removeObjectAtIndex:index];
        }
        obj
    }

    fn remove_last_object(&mut self) -> Id<Self::Item, Self::Own> {
        let obj = self.last_object().map(|obj| unsafe {
            Id::from_ptr(obj as *const _ as *mut Self::Item)
        });
        unsafe {
            let _: () = msg_send![self, removeLastObject];
        }
        // removeLastObject would have failed if the array is empty,
        // so we know this won't be None
        obj.unwrap()
    }

    fn remove_all_objects(&mut self) {
        unsafe {
            let _: () = msg_send![self, removeAllObjects];
        }
    }

    fn sort_by<F: FnMut(&Self::Item, &Self::Item) -> Ordering>(&mut self, compare: F) {
        extern fn compare_with_closure<T, F: FnMut(&T, &T) -> Ordering>(
                obj1: &T, obj2: &T, compare: &mut F) -> NSComparisonResult {
            NSComparisonResult::from_ordering((*compare)(obj1, obj2))
        }

        let mut closure = compare;
        unsafe {
            let _: () = msg_send![self, sortUsingFunction:compare_with_closure::<Self::Item, F>
                                                  context:&mut closure];
        }
    }
}

#[allow(missing_copy_implementations)]
pub enum NSMutableArray<T, O = Owned> { }

object_impl!(NSMutableArray<T, O>);

impl<T: INSObject, O: Ownership> INSArray for NSMutableArray<T, O> {
    type Item = T;
    type Own = O;
}

impl<T: INSObject> INSOwnedArray for NSMutableArray<T, Owned> { }

impl<T: INSObject> INSSharedArray for NSMutableArray<T, Shared> { }

impl<T: INSObject, O: Ownership> INSMutableArray for NSMutableArray<T, O> { }

impl<T> INSCopying for NSMutableArray<T, Shared> {
    type Output = NSSharedArray<T>;
}

impl<T> INSMutableCopying for NSMutableArray<T, Shared> {
    type Output = NSMutableSharedArray<T>;
}

impl<T: INSObject, O: Ownership> Index<usize> for NSMutableArray<T, O> {
    type Output = T;

    fn index(&self, index: &usize) -> &T {
        self.object_at(*index)
    }
}

pub type NSMutableSharedArray<T> = NSMutableArray<T, Shared>;

#[cfg(test)]
mod tests {
    use objc::{Id};
    use {INSObject, INSString, NSObject, NSString};
    use super::{INSArray, INSMutableArray, NSArray, NSMutableArray};

    fn sample_array(len: usize) -> Id<NSArray<NSObject>> {
        let mut vec: Vec<Id<NSObject>> = Vec::with_capacity(len);
        for _ in 0..len {
            vec.push(INSObject::new());
        }
        INSArray::from_vec(vec)
    }

    #[test]
    fn test_count() {
        let empty_array: Id<NSArray<NSObject>> = INSObject::new();
        assert!(empty_array.count() == 0);

        let array = sample_array(4);
        assert!(array.count() == 4);
    }

    #[test]
    fn test_object_at() {
        let array = sample_array(4);
        assert!(array.object_at(0) != array.object_at(3));
        assert!(array.first_object().unwrap() == array.object_at(0));
        assert!(array.last_object().unwrap() == array.object_at(3));

        let empty_array: Id<NSArray<NSObject>> = INSObject::new();
        assert!(empty_array.first_object().is_none());
        assert!(empty_array.last_object().is_none());
    }

    #[test]
    fn test_object_enumerator() {
        let array = sample_array(4);

        assert!(array.object_enumerator().count() == 4);
        assert!(array.object_enumerator()
                     .enumerate()
                     .all(|(i, obj)| obj == array.object_at(i)));
    }

    #[test]
    fn test_objects_in_range() {
        let array = sample_array(4);

        let middle_objs = array.objects_in_range(1..3);
        assert!(middle_objs.len() == 2);
        assert!(middle_objs[0] == array.object_at(1));
        assert!(middle_objs[1] == array.object_at(2));

        let empty_objs = array.objects_in_range(1..1);
        assert!(empty_objs.len() == 0);

        let all_objs = array.objects_in_range(0..4);
        assert!(all_objs.len() == 4);
    }

    #[test]
    fn test_into_vec() {
        let array = sample_array(4);

        let vec = INSArray::into_vec(array);
        assert!(vec.len() == 4);
    }

    #[test]
    fn test_add_object() {
        let mut array: Id<NSMutableArray<NSObject>> = INSObject::new();
        let obj: Id<NSObject> = INSObject::new();
        array.add_object(obj);

        assert!(array.count() == 1);
        assert!(array.object_at(0) == array.object_at(0));

        let obj: Id<NSObject> = INSObject::new();
        array.insert_object_at(0, obj);
        assert!(array.count() == 2);
    }

    #[test]
    fn test_replace_object() {
        let mut array: Id<NSMutableArray<NSObject>> = INSObject::new();
        let obj: Id<NSObject> = INSObject::new();
        array.add_object(obj);

        let obj: Id<NSObject> = INSObject::new();
        let old_obj = array.replace_object_at(0, obj);
        assert!(&*old_obj != array.object_at(0));
    }

    #[test]
    fn test_remove_object() {
        let mut array: Id<NSMutableArray<NSObject>> = INSObject::new();
        for _ in 0..4 {
            let obj: Id<NSObject> = INSObject::new();
            array.add_object(obj);
        }

        array.remove_object_at(1);
        assert!(array.count() == 3);

        array.remove_last_object();
        assert!(array.count() == 2);

        array.remove_all_objects();
        assert!(array.count() == 0);
    }

    #[test]
    fn test_sort() {
        let strings: Vec<Id<NSString>> = vec![
            INSString::from_str("hello"),
            INSString::from_str("hi"),
        ];
        let mut strings: Id<NSMutableArray<_>> = INSArray::from_vec(strings);

        strings.sort_by(|s1, s2| s1.as_str().len().cmp(&s2.as_str().len()));
        assert!(strings[0].as_str() == "hi");
        assert!(strings[1].as_str() == "hello");
    }
}
<|endoftext|>


<commit_before>#![allow(clippy::nonstandard_macro_braces)] // clippy bug, see https://github.com/rust-lang/rust-clippy/issues/7434

use anyhow::Result;
use log::debug;
use thiserror::Error;

pub mod imports;
pub mod rsrc;

use crate::{
    arch::Arch,
    aspace::RelativeAddressSpace,
    module::{Module, Permissions, Section},
    util, RVA, VA,
};

#[derive(Error, Debug)]
pub enum PEError {
    #[error("format not supported: {0}")]
    FormatNotSupported(String),

    #[error("malformed PE file: {0}")]
    MalformedPEFile(String),
}

// ref: https://docs.microsoft.com/en-us/windows/win32/api/dbghelp/nf-dbghelp-imagedirectoryentrytodata#parameters
pub const IMAGE_DIRECTORY_ENTRY_EXPORT: usize = 0;
pub const IMAGE_DIRECTORY_ENTRY_IMPORT: usize = 1;
pub const IMAGE_DIRECTORY_ENTRY_RESOURCE: usize = 2;
pub const IMAGE_DIRECTORY_ENTRY_EXCEPTION: usize = 3;
pub const IMAGE_DIRECTORY_ENTRY_SECURITY: usize = 4;
pub const IMAGE_DIRECTORY_ENTRY_BASERELOC: usize = 5;
pub const IMAGE_DIRECTORY_ENTRY_DEBUG: usize = 6;
pub const IMAGE_DIRECTORY_ENTRY_ARCHITECTURE: usize = 7;
pub const IMAGE_DIRECTORY_ENTRY_GLOBALPTR: usize = 8;
pub const IMAGE_DIRECTORY_ENTRY_TLS: usize = 9;
pub const IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG: usize = 10;
pub const IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT: usize = 11;
pub const IMAGE_DIRECTORY_ENTRY_IAT: usize = 12;
pub const IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT: usize = 13;
pub const IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR: usize = 14;
pub const IMAGE_DIRECTORY_MAX: usize = IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR;

pub struct DataDirectory {
    pub address: VA,
    pub size:    RVA,
}

/// A parsed and loaded PE file.
/// The `buf` field contains the raw data.
/// The `module` field contains an address space as the PE would be loaded.
pub struct PE {
    pub buf:    Vec<u8>,
    pub module: Module,
    pub header: goblin::pe::header::Header,
}

impl PE {
    pub fn from_bytes(buf: &[u8]) -> Result<PE> {
        load_pe(buf)
    }

    pub fn executable_sections<'b>(&'b self) -> Box<dyn Iterator<Item = &Section> + 'b> {
        Box::new(
            self.module
                .sections
                .iter()
                .filter(|section| section.permissions.intersects(Permissions::X)),
        )
    }

    pub fn pe(&self) -> Result<goblin::pe::PE> {
        get_pe(&self.buf)
    }

    pub fn get_data_directory(&self, data_directory: usize) -> Result<Option<DataDirectory>> {
        assert!(data_directory <= IMAGE_DIRECTORY_MAX);

        let opt_header = match self.header.optional_header {
            Some(opt_header) => opt_header,
            _ => return Ok(None),
        };

        match opt_header.data_directories.data_directories[data_directory] {
            Some(directory) => Ok(Some(DataDirectory {
                address: self.module.address_space.base_address + directory.virtual_address as VA,
                size:    directory.size as RVA,
            })),
            _ => Ok(None),
        }
    }
}

fn get_pe(buf: &[u8]) -> Result<goblin::pe::PE> {
    let pe = goblin::pe::PE::parse(buf)?;
    if let Some(opt) = pe.header.optional_header {
        if opt.data_directories.get_clr_runtime_header().is_some() {
            return Err(PEError::FormatNotSupported(".NET assembly".to_string()).into());
        }
    }
    Ok(pe)
}

#[allow(clippy::unnecessary_wraps)]
fn load_pe_header(buf: &[u8], pe: &goblin::pe::PE, base_address: VA) -> Result<Section> {
    let hdr_raw_size = match pe.header.optional_header {
        Some(opt) => opt.windows_fields.size_of_headers,
        // assumption: header is at most 0x200 bytes.
        _ => 0x200,
    };

    //   on disk:
    //
    //   +---------------------------------+
    //   |   header        |  sections...  |
    //   +---------------------------------+
    //   .                  \
    //   .  in memory:       \
    //   .                    \
    //   +-----------------+---+        +-------------
    //   |   header        |   |        |  sections...
    //   +-----------------+---+        +-------------
    //                     ^   ^
    //                     |   +--- virtual size
    //                     |        aligned to 0x200
    //                     +-- raw size
    //                         no alignment

    let hdr_raw_size = std::cmp::min(hdr_raw_size as usize, buf.len());
    let hdr_virt_size = util::align(hdr_raw_size as u64, 0x200);

    Ok(Section {
        physical_range: std::ops::Range {
            start: 0x0,
            end:   hdr_raw_size as u64,
        },
        virtual_range:  std::ops::Range {
            start: base_address,
            end:   base_address + hdr_virt_size,
        },
        permissions:    Permissions::R,
        name:           "header".to_string(),
    })
}

/// The section can be executed as code.
const IMAGE_SCN_MEM_EXECUTE: u32 = 0x2000_0000;

/// The section can be read.
const IMAGE_SCN_MEM_READ: u32 = 0x4000_0000;

/// The section can be written to.
const IMAGE_SCN_MEM_WRITE: u32 = 0x8000_0000;

#[allow(clippy::unnecessary_wraps)]
fn load_pe_section(
    base_address: VA,
    section_alignment: u64,
    section: &goblin::pe::section_table::SectionTable,
) -> Result<Section> {
    let name = String::from_utf8_lossy(&section.name[..])
        .into_owned()
        .trim_end_matches('\u{0}')
        .trim_end()
        .splitn(2, '\u{0}')
        .next()
        .unwrap()
        .to_string();

    let virtual_size = util::align(section.virtual_size as u64, section_alignment);

    let mut perms = Permissions::empty();
    if section.characteristics & IMAGE_SCN_MEM_READ > 0 {
        perms.insert(Permissions::R);
    }
    if section.characteristics & IMAGE_SCN_MEM_WRITE > 0 {
        perms.insert(Permissions::W);
    }
    if section.characteristics & IMAGE_SCN_MEM_EXECUTE > 0 {
        perms.insert(Permissions::X);
    }

    debug!(
        "pe: section: {} at {:#x}",
        name,
        base_address + section.virtual_address as u64
    );

    Ok(Section {
        physical_range: std::ops::Range {
            start: section.pointer_to_raw_data as u64,
            end:   (section.pointer_to_raw_data + section.size_of_raw_data) as u64,
        },
        virtual_range: std::ops::Range {
            start: base_address + section.virtual_address as u64,
            end:   base_address + section.virtual_address as u64 + virtual_size,
        },
        permissions: perms,
        name,
    })
}

// lots of further detail here: https://github.com/corkami/docs/blob/master/PE/PE.md
fn load_pe(buf: &[u8]) -> Result<PE> {
    let pe = get_pe(buf)?;

    let arch = match pe.is_64 {
        false => Arch::X32,
        true => Arch::X64,
    };
    debug!("pe: arch: {:?}", arch);

    let (base_address, section_alignment) = match pe.header.optional_header {
        Some(opt) => (
            opt.windows_fields.image_base,
            opt.windows_fields.section_alignment as u64,
        ),
        _ => {
            debug!("pe: base address: using default: 0x40:000");
            (0x40_0000, 0x1000)
        }
    };
    debug!("pe: base address: {:#x}", base_address);

    let mut sections = vec![load_pe_header(buf, &pe, base_address)?];
    for section in pe.sections.iter() {
        sections.push(load_pe_section(base_address, section_alignment as u64, section)?);
    }

    let max_address = sections.iter().map(|sec| sec.virtual_range.end).max().unwrap();
    let max_page_address = util::align(max_address, 0x1000) - base_address;
    debug!("pe: address space: capacity: {:#x}", max_page_address);

    let mut address_space = RelativeAddressSpace::with_capacity(max_page_address);

    for section in sections.iter() {
        let pstart = section.physical_range.start as usize;
        let pend = section.physical_range.end as usize;
        let psize = pend - pstart;
        let pbuf = &buf[pstart..pend];

        // the section range contains VAs,
        // while we're writing to the RelativeAddressSpace.
        // so shift down by `base_address`.
        let vstart = section.virtual_range.start;
        let rstart = vstart - base_address;
        let vsize = util::align(
            section.virtual_range.end - section.virtual_range.start,
            section_alignment,
        );
        let vend = vstart + vsize;
        let mut vbuf = vec![0u8; vsize as usize];

        if vsize as usize >= psize {
            // vsize > psize, so there will be NULL bytes padding the physical data.
            let dest = &mut vbuf[0..psize as usize];
            dest.copy_from_slice(pbuf);
        } else {
            // psize > vsize, but vsize wins, so we only read a subset of physical data.
            let src = &pbuf[0..vsize as usize];
            vbuf.copy_from_slice(src);
        }

        address_space.map.writezx(rstart, &vbuf)?;

        debug!(
            "pe: address space: mapped {:#x} - {:#x} {:?}",
            vstart, vend, section.permissions
        );
    }

    let module = Module {
        arch,
        sections,
        address_space: address_space.into_absolute(base_address)?,
    };

    debug!("pe: loaded");
    Ok(PE {
        buf: buf.to_vec(),
        module,
        header: pe.header,
    })
}

#[cfg(test)]
mod tests {
    use anyhow::Result;

    use crate::{aspace::AddressSpace, rsrc::*};

    #[test]
    fn base_address() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        assert_eq!(0x1_8000_0000, pe.module.address_space.base_address);

        Ok(())
    }

    #[test]
    fn mz_header() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        // relative read
        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        // absolute read
        assert_eq!(0x4d, pe.module.address_space.read_u8(0x1_8000_0000 + 0x0)?);
        assert_eq!(0x5a, pe.module.address_space.read_u8(0x1_8000_0000 + 0x1)?);

        Ok(())
    }

    #[test]
    fn k32() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    #[test]
    fn tiny() -> Result<()> {
        let buf = get_buf(Rsrc::TINY);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    #[test]
    fn nop() -> Result<()> {
        let buf = get_buf(Rsrc::NOP);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    #[test]
    fn mimi() -> Result<()> {
        let buf = get_buf(Rsrc::MIMI);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    // this demonstrates that the PE will be loaded and sections padded out to their
    // virtual range.
    #[test]
    fn read_each_section() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        for section in pe.module.sections.iter() {
            let start = section.virtual_range.start;
            let size = section.virtual_range.end - section.virtual_range.start;
            pe.module
                .address_space
                .read_bytes(start, size as usize)
                .expect(&format!("read section {} {:#x} {:#x}", section.name, start, size));
        }

        Ok(())
    }
}
<commit_msg>clippy<commit_after>#![allow(clippy::nonstandard_macro_braces)] // clippy bug, see https://github.com/rust-lang/rust-clippy/issues/7434

use anyhow::Result;
use log::debug;
use thiserror::Error;

pub mod imports;
pub mod rsrc;

use crate::{
    arch::Arch,
    aspace::RelativeAddressSpace,
    module::{Module, Permissions, Section},
    util, RVA, VA,
};

#[derive(Error, Debug)]
pub enum PEError {
    #[error("format not supported: {0}")]
    FormatNotSupported(String),

    #[error("malformed PE file: {0}")]
    MalformedPEFile(String),
}

// ref: https://docs.microsoft.com/en-us/windows/win32/api/dbghelp/nf-dbghelp-imagedirectoryentrytodata#parameters
pub const IMAGE_DIRECTORY_ENTRY_EXPORT: usize = 0;
pub const IMAGE_DIRECTORY_ENTRY_IMPORT: usize = 1;
pub const IMAGE_DIRECTORY_ENTRY_RESOURCE: usize = 2;
pub const IMAGE_DIRECTORY_ENTRY_EXCEPTION: usize = 3;
pub const IMAGE_DIRECTORY_ENTRY_SECURITY: usize = 4;
pub const IMAGE_DIRECTORY_ENTRY_BASERELOC: usize = 5;
pub const IMAGE_DIRECTORY_ENTRY_DEBUG: usize = 6;
pub const IMAGE_DIRECTORY_ENTRY_ARCHITECTURE: usize = 7;
pub const IMAGE_DIRECTORY_ENTRY_GLOBALPTR: usize = 8;
pub const IMAGE_DIRECTORY_ENTRY_TLS: usize = 9;
pub const IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG: usize = 10;
pub const IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT: usize = 11;
pub const IMAGE_DIRECTORY_ENTRY_IAT: usize = 12;
pub const IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT: usize = 13;
pub const IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR: usize = 14;
pub const IMAGE_DIRECTORY_MAX: usize = IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR;

pub struct DataDirectory {
    pub address: VA,
    pub size:    RVA,
}

/// A parsed and loaded PE file.
/// The `buf` field contains the raw data.
/// The `module` field contains an address space as the PE would be loaded.
pub struct PE {
    pub buf:    Vec<u8>,
    pub module: Module,
    pub header: goblin::pe::header::Header,
}

impl PE {
    pub fn from_bytes(buf: &[u8]) -> Result<PE> {
        load_pe(buf)
    }

    pub fn executable_sections<'b>(&'b self) -> Box<dyn Iterator<Item = &Section> + 'b> {
        Box::new(
            self.module
                .sections
                .iter()
                .filter(|section| section.permissions.intersects(Permissions::X)),
        )
    }

    pub fn pe(&self) -> Result<goblin::pe::PE> {
        get_pe(&self.buf)
    }

    pub fn get_data_directory(&self, data_directory: usize) -> Result<Option<DataDirectory>> {
        assert!(data_directory <= IMAGE_DIRECTORY_MAX);

        let opt_header = match self.header.optional_header {
            Some(opt_header) => opt_header,
            _ => return Ok(None),
        };

        match opt_header.data_directories.data_directories[data_directory] {
            Some(directory) => Ok(Some(DataDirectory {
                address: self.module.address_space.base_address + directory.virtual_address as VA,
                size:    directory.size as RVA,
            })),
            _ => Ok(None),
        }
    }
}

fn get_pe(buf: &[u8]) -> Result<goblin::pe::PE> {
    let pe = goblin::pe::PE::parse(buf)?;
    if let Some(opt) = pe.header.optional_header {
        if opt.data_directories.get_clr_runtime_header().is_some() {
            return Err(PEError::FormatNotSupported(".NET assembly".to_string()).into());
        }
    }
    Ok(pe)
}

#[allow(clippy::unnecessary_wraps)]
fn load_pe_header(buf: &[u8], pe: &goblin::pe::PE, base_address: VA) -> Result<Section> {
    let hdr_raw_size = match pe.header.optional_header {
        Some(opt) => opt.windows_fields.size_of_headers,
        // assumption: header is at most 0x200 bytes.
        _ => 0x200,
    };

    //   on disk:
    //
    //   +---------------------------------+
    //   |   header        |  sections...  |
    //   +---------------------------------+
    //   .                  \
    //   .  in memory:       \
    //   .                    \
    //   +-----------------+---+        +-------------
    //   |   header        |   |        |  sections...
    //   +-----------------+---+        +-------------
    //                     ^   ^
    //                     |   +--- virtual size
    //                     |        aligned to 0x200
    //                     +-- raw size
    //                         no alignment

    let hdr_raw_size = std::cmp::min(hdr_raw_size as usize, buf.len());
    let hdr_virt_size = util::align(hdr_raw_size as u64, 0x200);

    Ok(Section {
        physical_range: std::ops::Range {
            start: 0x0,
            end:   hdr_raw_size as u64,
        },
        virtual_range:  std::ops::Range {
            start: base_address,
            end:   base_address + hdr_virt_size,
        },
        permissions:    Permissions::R,
        name:           "header".to_string(),
    })
}

/// The section can be executed as code.
const IMAGE_SCN_MEM_EXECUTE: u32 = 0x2000_0000;

/// The section can be read.
const IMAGE_SCN_MEM_READ: u32 = 0x4000_0000;

/// The section can be written to.
const IMAGE_SCN_MEM_WRITE: u32 = 0x8000_0000;

#[allow(clippy::unnecessary_wraps)]
fn load_pe_section(
    base_address: VA,
    section_alignment: u64,
    section: &goblin::pe::section_table::SectionTable,
) -> Result<Section> {
    let section_name = String::from_utf8_lossy(&section.name[..]).into_owned();

    let trimmed_name = section_name
        .trim_end_matches('\u{0}')
        .trim_end();

    let name = trimmed_name
        .split_once('\u{0}')
        .map(|(name, _)| name)
        .unwrap_or_else(|| trimmed_name)
        .to_string();

    let virtual_size = util::align(section.virtual_size as u64, section_alignment);

    let mut perms = Permissions::empty();
    if section.characteristics & IMAGE_SCN_MEM_READ > 0 {
        perms.insert(Permissions::R);
    }
    if section.characteristics & IMAGE_SCN_MEM_WRITE > 0 {
        perms.insert(Permissions::W);
    }
    if section.characteristics & IMAGE_SCN_MEM_EXECUTE > 0 {
        perms.insert(Permissions::X);
    }

    debug!(
        "pe: section: {} at {:#x}",
        name,
        base_address + section.virtual_address as u64
    );

    Ok(Section {
        physical_range: std::ops::Range {
            start: section.pointer_to_raw_data as u64,
            end:   (section.pointer_to_raw_data + section.size_of_raw_data) as u64,
        },
        virtual_range: std::ops::Range {
            start: base_address + section.virtual_address as u64,
            end:   base_address + section.virtual_address as u64 + virtual_size,
        },
        permissions: perms,
        name,
    })
}

// lots of further detail here: https://github.com/corkami/docs/blob/master/PE/PE.md
fn load_pe(buf: &[u8]) -> Result<PE> {
    let pe = get_pe(buf)?;

    let arch = match pe.is_64 {
        false => Arch::X32,
        true => Arch::X64,
    };
    debug!("pe: arch: {:?}", arch);

    let (base_address, section_alignment) = match pe.header.optional_header {
        Some(opt) => (
            opt.windows_fields.image_base,
            opt.windows_fields.section_alignment as u64,
        ),
        _ => {
            debug!("pe: base address: using default: 0x40:000");
            (0x40_0000, 0x1000)
        }
    };
    debug!("pe: base address: {:#x}", base_address);

    let mut sections = vec![load_pe_header(buf, &pe, base_address)?];
    for section in pe.sections.iter() {
        sections.push(load_pe_section(base_address, section_alignment as u64, section)?);
    }

    let max_address = sections.iter().map(|sec| sec.virtual_range.end).max().unwrap();
    let max_page_address = util::align(max_address, 0x1000) - base_address;
    debug!("pe: address space: capacity: {:#x}", max_page_address);

    let mut address_space = RelativeAddressSpace::with_capacity(max_page_address);

    for section in sections.iter() {
        let pstart = section.physical_range.start as usize;
        let pend = section.physical_range.end as usize;
        let psize = pend - pstart;
        let pbuf = &buf[pstart..pend];

        // the section range contains VAs,
        // while we're writing to the RelativeAddressSpace.
        // so shift down by `base_address`.
        let vstart = section.virtual_range.start;
        let rstart = vstart - base_address;
        let vsize = util::align(
            section.virtual_range.end - section.virtual_range.start,
            section_alignment,
        );
        let vend = vstart + vsize;
        let mut vbuf = vec![0u8; vsize as usize];

        if vsize as usize >= psize {
            // vsize > psize, so there will be NULL bytes padding the physical data.
            let dest = &mut vbuf[0..psize as usize];
            dest.copy_from_slice(pbuf);
        } else {
            // psize > vsize, but vsize wins, so we only read a subset of physical data.
            let src = &pbuf[0..vsize as usize];
            vbuf.copy_from_slice(src);
        }

        address_space.map.writezx(rstart, &vbuf)?;

        debug!(
            "pe: address space: mapped {:#x} - {:#x} {:?}",
            vstart, vend, section.permissions
        );
    }

    let module = Module {
        arch,
        sections,
        address_space: address_space.into_absolute(base_address)?,
    };

    debug!("pe: loaded");
    Ok(PE {
        buf: buf.to_vec(),
        module,
        header: pe.header,
    })
}

#[cfg(test)]
mod tests {
    use anyhow::Result;

    use crate::{aspace::AddressSpace, rsrc::*};

    #[test]
    fn base_address() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        assert_eq!(0x1_8000_0000, pe.module.address_space.base_address);

        Ok(())
    }

    #[test]
    fn mz_header() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        // relative read
        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        // absolute read
        assert_eq!(0x4d, pe.module.address_space.read_u8(0x1_8000_0000 + 0x0)?);
        assert_eq!(0x5a, pe.module.address_space.read_u8(0x1_8000_0000 + 0x1)?);

        Ok(())
    }

    #[test]
    fn k32() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    #[test]
    fn tiny() -> Result<()> {
        let buf = get_buf(Rsrc::TINY);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    #[test]
    fn nop() -> Result<()> {
        let buf = get_buf(Rsrc::NOP);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    #[test]
    fn mimi() -> Result<()> {
        let buf = get_buf(Rsrc::MIMI);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        assert_eq!(0x4d, pe.module.address_space.relative.read_u8(0x0)?);
        assert_eq!(0x5a, pe.module.address_space.relative.read_u8(0x1)?);

        Ok(())
    }

    // this demonstrates that the PE will be loaded and sections padded out to their
    // virtual range.
    #[test]
    fn read_each_section() -> Result<()> {
        let buf = get_buf(Rsrc::K32);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        for section in pe.module.sections.iter() {
            let start = section.virtual_range.start;
            let size = section.virtual_range.end - section.virtual_range.start;
            pe.module
                .address_space
                .read_bytes(start, size as usize)
                .expect(&format!("read section {} {:#x} {:#x}", section.name, start, size));
        }

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


<commit_before>include!("header.rs");

#[test]
fn ssqrt_symmetric_random() {
    let r_dist = RealNormal::new(0., 1.);
    let mut a = Array::<f64, _>::random((3, 3), r_dist);
    a = a.dot(&a.t());
    let ar = a.clone().ssqrt().unwrap();
    all_close_l2(&ar.clone().reversed_axes(), &ar, 1e-7).unwrap();
}

#[test]
fn ssqrt_symmetric_random_t() {
    let r_dist = RealNormal::new(0., 1.);
    let mut a = Array::<f64, _>::random((3, 3), r_dist);
    a = a.dot(&a.t()).reversed_axes();
    let ar = a.clone().ssqrt().unwrap();
    all_close_l2(&ar.clone().reversed_axes(), &ar, 1e-7).unwrap();
}

#[test]
fn ssqrt_sqrt_random() {
    let r_dist = RealNormal::new(0., 1.);
    let mut a = Array::<f64, _>::random((3, 3), r_dist);
    a = a.dot(&a.t());
    let ar = a.clone().ssqrt().unwrap();
    all_close_l2(&ar.clone().reversed_axes(), &ar, 1e-7).unwrap();
}

#[test]
fn ssqrt_sqrt_random_t() {
    let r_dist = RealNormal::new(0., 1.);
    let mut a = Array::<f64, _>::random((3, 3), r_dist);
    a = a.dot(&a.t()).reversed_axes();
    let ar = a.clone().ssqrt().unwrap();
    all_close_l2(&ar.clone().reversed_axes(), &ar, 1e-7).unwrap();
}
<commit_msg>Test for ssqrt<commit_after>include!("header.rs");

macro_rules! impl_test{
    ($modname:ident, $clone:ident) => {
mod $modname {
    use super::random_hermite;
    use ndarray_linalg::prelude::*;
    #[test]
    fn ssqrt() {
        let a = random_hermite(3);
        let ar = a.$clone().ssqrt().unwrap();
        all_close_l2(&ar.clone().t(), &ar, 1e-7).expect("not symmetric");
        all_close_l2(&ar.dot(&ar), &a, 1e-7).expect("not sqrt");
    }
    #[test]
    fn ssqrt_t() {
        let a = random_hermite(3).reversed_axes();
        let ar = a.$clone().ssqrt().unwrap();
        all_close_l2(&ar.clone().t(), &ar, 1e-7).expect("not symmetric");
        all_close_l2(&ar.dot(&ar), &a, 1e-7).expect("not sqrt");
    }
}
}} // impl_test

impl_test!(owned, clone);
impl_test!(shared, to_shared);
<|endoftext|>