<commit_before>use ffi::prelude::LLVMTypeRef;
use ffi::{core, target, LLVMTypeKind};
use libc::{c_int, c_uint};
use compile::Compile;
use context::{Context, GetContext};
use target::TargetData;
use util::{self, Sub};
use std::{fmt, mem};
use std::marker::PhantomData;
use std::iter::Iterator;
use std::ops::Deref;

macro_rules! sub {
    ($ty:ty, $kind:ident) => (
unsafe impl Sub<::Type> for $ty {
    fn is(ty: &Type) -> bool {
        unsafe {
            let kind = core::LLVMGetTypeKind(ty.into());
            kind as c_uint == LLVMTypeKind::$kind as c_uint
        }
    }
}
deref!{$ty, Type}
    )
}

/// Defines how a value should be laid out in memory.
pub struct Type(PhantomData<[u8]>);
native_ref!(&Type = LLVMTypeRef);
impl Type {
    #[inline(always)]
    /// Get the type given as an LLVM type descriptor in the context given.
    pub fn get<'a, T>(context:&'a Context) -> &'a Type where T:Compile<'a> {
        T::get_type(context)
    }
    /// Make a new array with the length given.
    pub fn new_array<'a>(element: &'a Type, length: usize) -> &'a Type {
        unsafe { core::LLVMArrayType(element.into(), length as c_uint) }.into()
    }
    /// Make a new vector with the length given.
    pub fn new_vector<'a>(element: &'a Type, length: usize) -> &'a Type {
        unsafe { core::LLVMVectorType(element.into(), length as c_uint) }.into()
    }
    /// Returns true if the size of the type is known at compile-time.
    ///
    /// This is equivalent to the type implementing `Sized` in Rust
    pub fn is_sized(&self) -> bool {
        unsafe { core::LLVMTypeIsSized(self.into()) != 0 }
    }
    /// Returns true if this type is a function.
    pub fn is_function(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMFunctionTypeKind as c_uint
    }
    /// Returns true if this type is a struct.
    pub fn is_struct(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMStructTypeKind as c_uint
    }
    /// Returns true if this type is void.
    pub fn is_void(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMVoidTypeKind as c_uint
    }
    /// Returns true if this type is a pointer.
    pub fn is_pointer(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMPointerTypeKind as c_uint
    }
    /// Returns true if this type is an integer.
    pub fn is_integer(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMIntegerTypeKind as c_uint
    }
    /// Returns true if this type is any floating-point number.
    pub fn is_float(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) } as c_uint;
        kind == LLVMTypeKind::LLVMHalfTypeKind as c_uint ||
        kind == LLVMTypeKind::LLVMFloatTypeKind as c_uint ||
        kind == LLVMTypeKind::LLVMDoubleTypeKind as c_uint
    }
    /// Returns the size of the type in bytes.
    pub fn get_size(&self, target: &TargetData) -> usize {
        unsafe { target::LLVMABISizeOfType(target.into(), self.into()) as usize }
    }
}
get_context!(Type, LLVMGetTypeContext);
to_str!(Type, LLVMPrintTypeToString);

/// A structure type, such as a tuple or struct.
pub struct StructType;
native_ref!(&StructType = LLVMTypeRef);
sub!{StructType, LLVMStructTypeKind}
impl StructType {
    /// Make a new struct with the given fields and packed representation.
    pub fn new<'a>(context: &'a Context, fields: &[&'a Type], packed: bool) -> &'a StructType {
        unsafe { core::LLVMStructTypeInContext(context.into(), fields.as_ptr() as *mut LLVMTypeRef, fields.len() as c_uint, packed as c_int) }.into()
    }
    /// Make a new named struct with the given fields and packed representation.
    pub fn new_named<'a>(context: &'a Context, name: &str, fields: &[&'a Type], packed: bool) -> &'a StructType {
        util::with_cstr(name, |name| unsafe {
            let ty = core::LLVMStructCreateNamed(context.into(), name);
            core::LLVMStructSetBody(ty, fields.as_ptr() as *mut LLVMTypeRef, fields.len() as c_uint, packed as c_int);
            ty.into()
        })
    }
    /// Returns the elements that make up this struct.
    pub fn get_elements(&self) -> Vec<&Type> {
        unsafe {
            let size = core::LLVMCountStructElementTypes(self.into());
            let mut els:Vec<_> = (0..size).map(|_| mem::uninitialized()).collect();
            core::LLVMGetStructElementTypes(self.into(), els.as_mut_ptr() as *mut LLVMTypeRef);
            els
        }
    }
}
get_context!(StructType, LLVMGetTypeContext);
to_str!(StructType, LLVMPrintTypeToString);

/// A function signature type.
pub struct FunctionType;
native_ref!(&FunctionType = LLVMTypeRef);
deref!(FunctionType, Type);
unsafe impl Sub<Type> for FunctionType {
    fn is(mut ty: &Type) -> bool {
        unsafe {
            while let Some(ptr) = PointerType::from_super(ty) {
                ty = ptr.get_element();
            }
            let kind = core::LLVMGetTypeKind(ty.into());
            kind as c_uint == LLVMTypeKind::LLVMFunctionTypeKind as c_uint
        }
    }
}
impl FunctionType {
    /// Make a new function signature with the return type and arguments given.
    pub fn new<'a>(ret: &'a Type, args: &[&'a Type]) -> &'a FunctionType {
        unsafe { core::LLVMFunctionType(ret.into(), args.as_ptr() as *mut LLVMTypeRef, args.len() as c_uint, 0) }.into()
    }
    /// Returns the number of parameters this signature takes.
    pub fn num_params(&self) -> usize {
        unsafe { core::LLVMCountParamTypes(self.into()) as usize }
    }
    /// Returns a vector of this signature's parameters' types.
    pub fn get_params(&self) -> Vec<&Type> {
        unsafe {
            let count = core::LLVMCountParamTypes(self.into());
            let mut types:Vec<_> = (0..count).map(|_| mem::uninitialized()).collect();
            core::LLVMGetParamTypes(self.into(), types.as_mut_ptr() as *mut LLVMTypeRef);
            types
        }
    }
    /// Returns the type that this function returns.
    pub fn get_return(&self) -> &Type {
        unsafe { core::LLVMGetReturnType(self.into()).into() }
    }
}
get_context!(FunctionType, LLVMGetTypeContext);
to_str!(FunctionType, LLVMPrintTypeToString);

/// A pointer type.
pub struct PointerType;
sub!{PointerType, LLVMPointerTypeKind}
native_ref!(&PointerType = LLVMTypeRef);
impl PointerType {
    /// Make a new pointer type with the given element type.
    pub fn new(elem: &Type) -> &Type {
        unsafe { core::LLVMPointerType(elem.into(), 0 as c_uint) }.into()
    }
    /// Returns the element of this pointer type.
    pub fn get_element(&self) -> &Type {
        unsafe { mem::transmute(core::LLVMGetElementType(self.into())) }
    }
}
get_context!(PointerType, LLVMGetTypeContext);
to_str!(PointerType, LLVMPrintTypeToString);<commit_msg>Tidy up types module and add `IntegerType`<commit_after>use ffi::prelude::LLVMTypeRef;
use ffi::{core, target, LLVMTypeKind};
use libc::{c_int, c_uint};
use compile::Compile;
use context::{Context, GetContext};
use target::TargetData;
use util::{self, Sub};
use std::{fmt, mem};
use std::marker::PhantomData;
use std::iter::Iterator;
use std::ops::Deref;

macro_rules! sub {
    ($ty:ty, $kind:ident) => (
unsafe impl Sub<::Type> for $ty {
    fn is(ty: &Type) -> bool {
        unsafe {
            let kind = core::LLVMGetTypeKind(ty.into());
            kind as c_uint == LLVMTypeKind::$kind as c_uint
        }
    }
}
deref!{$ty, Type}
    )
}

/// Defines how a value should be laid out in memory.
pub struct Type(PhantomData<[u8]>);
native_ref!{&Type = LLVMTypeRef}
get_context!{Type, LLVMGetTypeContext}
to_str!{Type, LLVMPrintTypeToString}
impl Type {
    #[inline(always)]
    /// Get the type given as an LLVM type descriptor in the context given.
    pub fn get<'a, T>(context:&'a Context) -> &'a Type where T:Compile<'a> {
        T::get_type(context)
    }
    /// Make a new array with the length given.
    pub fn new_array<'a>(element: &'a Type, length: usize) -> &'a Type {
        unsafe { core::LLVMArrayType(element.into(), length as c_uint) }.into()
    }
    /// Make a new vector with the length given.
    pub fn new_vector<'a>(element: &'a Type, length: usize) -> &'a Type {
        unsafe { core::LLVMVectorType(element.into(), length as c_uint) }.into()
    }
    /// Returns true if the size of the type is known at compile-time.
    ///
    /// This is equivalent to the type implementing `Sized` in Rust
    pub fn is_sized(&self) -> bool {
        unsafe { core::LLVMTypeIsSized(self.into()) != 0 }
    }
    /// Returns true if this type is a function.
    ///
    /// This is equivalent to `FunctionType::is`.
    pub fn is_function(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMFunctionTypeKind as c_uint
    }
    /// Returns true if this type is a struct.
    ///
    /// This is equivalent to `StructType::is`.
    pub fn is_struct(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMStructTypeKind as c_uint
    }
    /// Returns true if this type is void.
    pub fn is_void(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMVoidTypeKind as c_uint
    }
    /// Returns true if this type is a pointer.
    ///
    /// This is equivalent to `PointerType::is`.
    pub fn is_pointer(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMPointerTypeKind as c_uint
    }
    /// Returns true if this type is an integer.
    pub fn is_integer(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) };
        kind as c_uint == LLVMTypeKind::LLVMIntegerTypeKind as c_uint
    }
    /// Returns true if this type is any floating-point number.
    pub fn is_float(&self) -> bool {
        let kind = unsafe { core::LLVMGetTypeKind(self.into()) } as c_uint;
        kind == LLVMTypeKind::LLVMHalfTypeKind as c_uint ||
        kind == LLVMTypeKind::LLVMFloatTypeKind as c_uint ||
        kind == LLVMTypeKind::LLVMDoubleTypeKind as c_uint
    }
    /// Returns the size of the type in bytes.
    pub fn get_size(&self, target: &TargetData) -> usize {
        unsafe { target::LLVMABISizeOfType(target.into(), self.into()) as usize }
    }
}

/// A structure type, such as a tuple or struct.
pub struct StructType;
native_ref!{&StructType = LLVMTypeRef}
get_context!{StructType, LLVMGetTypeContext}
to_str!{StructType, LLVMPrintTypeToString}
sub!{StructType, LLVMStructTypeKind}
impl StructType {
    /// Make a new struct with the given fields and packed representation.
    pub fn new<'a>(context: &'a Context, fields: &[&'a Type], packed: bool) -> &'a StructType {
        unsafe { core::LLVMStructTypeInContext(context.into(), fields.as_ptr() as *mut LLVMTypeRef, fields.len() as c_uint, packed as c_int) }.into()
    }
    /// Make a new named struct with the given fields and packed representation.
    pub fn new_named<'a>(context: &'a Context, name: &str, fields: &[&'a Type], packed: bool) -> &'a StructType {
        util::with_cstr(name, |name| unsafe {
            let ty = core::LLVMStructCreateNamed(context.into(), name);
            core::LLVMStructSetBody(ty, fields.as_ptr() as *mut LLVMTypeRef, fields.len() as c_uint, packed as c_int);
            ty.into()
        })
    }
    /// Returns the elements that make up this struct.
    pub fn get_elements(&self) -> Vec<&Type> {
        unsafe {
            let size = core::LLVMCountStructElementTypes(self.into());
            let mut els:Vec<_> = (0..size).map(|_| mem::uninitialized()).collect();
            core::LLVMGetStructElementTypes(self.into(), els.as_mut_ptr() as *mut LLVMTypeRef);
            els
        }
    }
}

/// A function signature type.
pub struct FunctionType;
native_ref!{&FunctionType = LLVMTypeRef}
get_context!{FunctionType, LLVMGetTypeContext}
to_str!{FunctionType, LLVMPrintTypeToString}
deref!{FunctionType, Type}
unsafe impl Sub<Type> for FunctionType {
    fn is(mut ty: &Type) -> bool {
        unsafe {
            while let Some(ptr) = PointerType::from_super(ty) {
                ty = ptr.get_element();
            }
            let kind = core::LLVMGetTypeKind(ty.into());
            kind as c_uint == LLVMTypeKind::LLVMFunctionTypeKind as c_uint
        }
    }
}
impl FunctionType {
    /// Make a new function signature with the return type and arguments given.
    pub fn new<'a>(ret: &'a Type, args: &[&'a Type]) -> &'a FunctionType {
        unsafe { core::LLVMFunctionType(ret.into(), args.as_ptr() as *mut LLVMTypeRef, args.len() as c_uint, 0) }.into()
    }
    /// Returns the number of parameters this signature takes.
    pub fn num_params(&self) -> usize {
        unsafe { core::LLVMCountParamTypes(self.into()) as usize }
    }
    /// Returns a vector of this signature's parameters' types.
    pub fn get_params(&self) -> Vec<&Type> {
        unsafe {
            let count = core::LLVMCountParamTypes(self.into());
            let mut types:Vec<_> = (0..count).map(|_| mem::uninitialized()).collect();
            core::LLVMGetParamTypes(self.into(), types.as_mut_ptr() as *mut LLVMTypeRef);
            types
        }
    }
    /// Returns the type that this function returns.
    pub fn get_return(&self) -> &Type {
        unsafe { core::LLVMGetReturnType(self.into()).into() }
    }
}

/// A pointer type.
pub struct PointerType;
native_ref!{&PointerType = LLVMTypeRef}
get_context!{PointerType, LLVMGetTypeContext}
to_str!{PointerType, LLVMPrintTypeToString}
sub!{PointerType, LLVMPointerTypeKind}
impl PointerType {
    /// Make a new pointer type with the given element type.
    pub fn new(elem: &Type) -> &Type {
        unsafe { core::LLVMPointerType(elem.into(), 0 as c_uint) }.into()
    }
    /// Returns the element of this pointer type.
    pub fn get_element(&self) -> &Type {
        unsafe { mem::transmute(core::LLVMGetElementType(self.into())) }
    }
}

/// An integer type.
pub struct IntegerType;
native_ref!{&IntegerType = LLVMTypeRef}
get_context!{IntegerType, LLVMGetTypeContext}
to_str!{IntegerType, LLVMPrintTypeToString}
sub!{IntegerType, LLVMPointerTypeKind}
impl IntegerType {
    /// Make a new integer type that will be the size of the given number of bits.
    pub fn new(context: &Context, numbits: usize) -> &IntegerType {
        unsafe { core::LLVMIntTypeInContext(context.into(), numbits as c_uint) }.into()
    }
    /// Returns how long an integer of this type is, in bits.
    pub fn get_width(&self) -> usize {
        unsafe { core::LLVMGetIntTypeWidth (self.into()) as usize }
    }
}<|endoftext|>


<commit_before>//! Support datetimes and timedeltas
//!
//! [The corresponding section][datetime] of the NumPy documentation contains more information.
//!
//! # Example
//!
//! ```
//! use numpy::{datetime::{units, Datetime, Timedelta}, PyArray1};
//! use pyo3::Python;
//! # use pyo3::types::PyDict;
//!
//! Python::with_gil(|py| {
//! #    let locals = py
//! #        .eval("{ 'np': __import__('numpy') }", None, None)
//! #        .unwrap()
//! #        .downcast::<PyDict>()
//! #        .unwrap();
//! #
//!     let array = py
//!         .eval(
//!             "np.array([np.datetime64('2017-04-21')])",
//!             None,
//!             Some(locals),
//!         )
//!         .unwrap()
//!         .downcast::<PyArray1<Datetime<units::Days>>>()
//!         .unwrap();
//!
//!     assert_eq!(
//!         array.get_owned(0).unwrap(),
//!         Datetime::<units::Days>::from(17_277)
//!     );
//!
//!     let array = py
//!         .eval(
//!             "np.array([np.datetime64('2022-03-29')]) - np.array([np.datetime64('2017-04-21')])",
//!             None,
//!             Some(locals),
//!         )
//!         .unwrap()
//!         .downcast::<PyArray1<Timedelta<units::Days>>>()
//!         .unwrap();
//!
//!     assert_eq!(
//!         array.get_owned(0).unwrap(),
//!         Timedelta::<units::Days>::from(1_803)
//!     );
//! });
//! ```
//!
//! [datetime]: https://numpy.org/doc/stable/reference/arrays.datetime.html

use std::cell::UnsafeCell;
use std::collections::hash_map::Entry;
use std::fmt;
use std::hash::Hash;
use std::marker::PhantomData;

use ahash::AHashMap;
use pyo3::{Py, Python};

use crate::dtype::{Element, PyArrayDescr};
use crate::npyffi::{PyArray_DatetimeDTypeMetaData, NPY_DATETIMEUNIT, NPY_TYPES};

/// Represents the [datetime units][datetime-units] supported by NumPy
///
/// [datetime-units]: https://numpy.org/doc/stable/reference/arrays.datetime.html#datetime-units
pub trait Unit: Send + Sync + Clone + Copy + PartialEq + Eq + Hash + PartialOrd + Ord {
    /// The matching NumPy [datetime unit code][NPY_DATETIMEUNIT]
    ///
    /// [NPY_DATETIMEUNIT]: https://github.com/numpy/numpy/blob/4c60b3263ac50e5e72f6a909e156314fc3c9cba0/numpy/core/include/numpy/ndarraytypes.h#L276
    const UNIT: NPY_DATETIMEUNIT;

    /// The abbrevation used for debug formatting
    const ABBREV: &'static str;
}

macro_rules! define_units {
    ($($(#[$meta:meta])* $struct:ident => $unit:ident $abbrev:literal,)+) => {
        $(

        $(#[$meta])*
        #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
        pub struct $struct;

        impl Unit for $struct {
            const UNIT: NPY_DATETIMEUNIT = NPY_DATETIMEUNIT::$unit;

            const ABBREV: &'static str = $abbrev;
        }

        )+
    };
}

/// Predefined implementors of the [`Unit`] trait
pub mod units {
    use super::*;

    define_units!(
        #[doc = "Years, i.e. 12 months"]
        Years => NPY_FR_Y "a",
        #[doc = "Months, i.e. 30 days"]
        Months => NPY_FR_M "mo",
        #[doc = "Weeks, i.e. 7 days"]
        Weeks => NPY_FR_W "w",
        #[doc = "Days, i.e. 24 hours"]
        Days => NPY_FR_D "d",
        #[doc = "Hours, i.e. 60 minutes"]
        Hours => NPY_FR_h "h",
        #[doc = "Minutes, i.e. 60 seconds"]
        Minutes => NPY_FR_m "min",
        #[doc = "Seconds"]
        Seconds => NPY_FR_s "s",
        #[doc = "Milliseconds, i.e. 10^-3 seconds"]
        Milliseconds => NPY_FR_ms "ms",
        #[doc = "Microseconds, i.e. 10^-6 seconds"]
        Microseconds => NPY_FR_us "µs",
        #[doc = "Nanoseconds, i.e. 10^-9 seconds"]
        Nanoseconds => NPY_FR_ns "ns",
        #[doc = "Picoseconds, i.e. 10^-12 seconds"]
        Picoseconds => NPY_FR_ps "ps",
        #[doc = "Femtoseconds, i.e. 10^-15 seconds"]
        Femtoseconds => NPY_FR_fs "fs",
        #[doc = "Attoseconds, i.e. 10^-18 seconds"]
        Attoseconds => NPY_FR_as "as",
    );
}

/// Corresponds to the [`datetime64`][scalars-datetime64] scalar type
///
/// [scalars-datetime64]: https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.datetime64
#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(transparent)]
pub struct Datetime(i64, PhantomData);

impl From<i64> for Datetime {
    fn from(val: i64) -> Self {
        Self(val, PhantomData)
    }
}

impl From<Datetime> for i64 {
    fn from(val: Datetime) -> Self {
        val.0
    }
}

unsafe impl Element for Datetime {
    const IS_COPY: bool = true;

    fn get_dtype(py: Python) -> &PyArrayDescr {
        static DTYPES: TypeDescriptors = unsafe { TypeDescriptors::new(NPY_TYPES::NPY_DATETIME) };

        DTYPES.from_unit(py, U::UNIT)
    }
}

impl fmt::Debug for Datetime {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Datetime({} {})", self.0, U::ABBREV)
    }
}

/// Corresponds to the [`timedelta64`][scalars-datetime64] scalar type
///
/// [scalars-timedelta64]: https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.timedelta64
#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(transparent)]
pub struct Timedelta(i64, PhantomData);

impl From<i64> for Timedelta {
    fn from(val: i64) -> Self {
        Self(val, PhantomData)
    }
}

impl From<Timedelta> for i64 {
    fn from(val: Timedelta) -> Self {
        val.0
    }
}

unsafe impl Element for Timedelta {
    const IS_COPY: bool = true;

    fn get_dtype(py: Python) -> &PyArrayDescr {
        static DTYPES: TypeDescriptors = unsafe { TypeDescriptors::new(NPY_TYPES::NPY_TIMEDELTA) };

        DTYPES.from_unit(py, U::UNIT)
    }
}

impl fmt::Debug for Timedelta {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Timedelta({} {})", self.0, U::ABBREV)
    }
}

struct TypeDescriptors {
    npy_type: NPY_TYPES,
    dtypes: UnsafeCell<Option<AHashMap<NPY_DATETIMEUNIT, Py<PyArrayDescr>>>>,
}

unsafe impl Sync for TypeDescriptors {}

impl TypeDescriptors {
    /// `npy_type` must be either `NPY_DATETIME` or `NPY_TIMEDELTA`.
    const unsafe fn new(npy_type: NPY_TYPES) -> Self {
        Self {
            npy_type,
            dtypes: UnsafeCell::new(None),
        }
    }

    #[allow(clippy::mut_from_ref)]
    unsafe fn get(&self) -> &mut AHashMap<NPY_DATETIMEUNIT, Py<PyArrayDescr>> {
        (*self.dtypes.get()).get_or_insert_with(AHashMap::new)
    }

    #[allow(clippy::wrong_self_convention)]
    fn from_unit<'py>(&'py self, py: Python<'py>, unit: NPY_DATETIMEUNIT) -> &'py PyArrayDescr {
        // SAFETY: We hold the GIL and we do not call into user code which might re-enter.
        let dtypes = unsafe { self.get() };

        match dtypes.entry(unit) {
            Entry::Occupied(entry) => entry.into_mut().as_ref(py),
            Entry::Vacant(entry) => {
                let dtype = PyArrayDescr::new_from_npy_type(py, self.npy_type);

                // SAFETY: `self.npy_type` is either `NPY_DATETIME` or `NPY_TIMEDELTA` which implies the type of `c_metadata`.
                unsafe {
                    let metadata = &mut *((*dtype.as_dtype_ptr()).c_metadata
                        as *mut PyArray_DatetimeDTypeMetaData);

                    metadata.meta.base = unit;
                    metadata.meta.num = 1;
                }

                entry.insert(dtype.into()).as_ref(py)
            }
        }
    }
}

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

    use pyo3::{
        py_run,
        types::{PyDict, PyModule},
    };

    use crate::array::PyArray1;

    #[test]
    fn from_python_to_rust() {
        Python::with_gil(|py| {
            let locals = py
                .eval("{ 'np': __import__('numpy') }", None, None)
                .unwrap()
                .downcast::<PyDict>()
                .unwrap();

            let array = py
                .eval(
                    "np.array([np.datetime64('1970-01-01')])",
                    None,
                    Some(locals),
                )
                .unwrap()
                .downcast::<PyArray1<Datetime<units::Days>>>()
                .unwrap();

            let value: i64 = array.get_owned(0).unwrap().into();
            assert_eq!(value, 0);
        });
    }

    #[test]
    fn from_rust_to_python() {
        Python::with_gil(|py| {
            let array = PyArray1::<Timedelta<units::Minutes>>::zeros(py, 1, false);

            *array.readwrite().get_mut(0).unwrap() = Timedelta::<units::Minutes>::from(5);

            let np = py
                .eval("__import__('numpy')", None, None)
                .unwrap()
                .downcast::<PyModule>()
                .unwrap();

            py_run!(py, array np, "assert array.dtype == np.dtype('timedelta64[m]')");
            py_run!(py, array np, "assert array[0] == np.timedelta64(5, 'm')");
        });
    }

    #[test]
    fn debug_formatting() {
        assert_eq!(
            format!("{:?}", Datetime::<units::Days>::from(28)),
            "Datetime(28 d)"
        );

        assert_eq!(
            format!("{:?}", Timedelta::<units::Milliseconds>::from(160)),
            "Timedelta(160 ms)"
        );
    }

    #[test]
    fn unit_conversion() {
        #[track_caller]
        fn convert<S: Unit, D: Unit>(py: Python<'_>, expected_value: i64) {
            let array = PyArray1::<Timedelta<S>>::from_slice(py, &[Timedelta::<S>::from(1)]);
            let array = array.cast::<Timedelta<D>>(false).unwrap();

            let value: i64 = array.get_owned(0).unwrap().into();
            assert_eq!(value, expected_value);
        }

        Python::with_gil(|py| {
            convert::<units::Years, units::Days>(py, (97 + 400 * 365) / 400);
            convert::<units::Months, units::Days>(py, (97 + 400 * 365) / 400 / 12);

            convert::<units::Weeks, units::Seconds>(py, 7 * 24 * 60 * 60);
            convert::<units::Days, units::Seconds>(py, 24 * 60 * 60);
            convert::<units::Hours, units::Seconds>(py, 60 * 60);
            convert::<units::Minutes, units::Seconds>(py, 60);

            convert::<units::Seconds, units::Milliseconds>(py, 1_000);
            convert::<units::Seconds, units::Microseconds>(py, 1_000_000);
            convert::<units::Seconds, units::Nanoseconds>(py, 1_000_000_000);
            convert::<units::Seconds, units::Picoseconds>(py, 1_000_000_000_000);
            convert::<units::Seconds, units::Femtoseconds>(py, 1_000_000_000_000_000);

            convert::<units::Femtoseconds, units::Attoseconds>(py, 1_000);
        });
    }
}
<commit_msg>Flesh out the datetime module docs a bit.<commit_after>//! Support datetimes and timedeltas
//!
//! This module provides wrappers for NumPy's [`datetime64`][scalars-datetime64] and [`timedelta64`][scalars-timedelta64] types
//! which are used for time keeping with with an emphasis on scientific applications.
//! This means that while these types differentiate absolute and relative quantities, they ignore calendars (a month is always 30.44 days) and time zones.
//! On the other hand, their flexible units enable them to support either a large range (up to 264 years) or high precision (down to 10-18 seconds).
//!
//! [The corresponding section][datetime] of the NumPy documentation contains more information.
//!
//! # Example
//!
//! ```
//! use numpy::{datetime::{units, Datetime, Timedelta}, PyArray1};
//! use pyo3::Python;
//! # use pyo3::types::PyDict;
//!
//! Python::with_gil(|py| {
//! #    let locals = py
//! #        .eval("{ 'np': __import__('numpy') }", None, None)
//! #        .unwrap()
//! #        .downcast::<PyDict>()
//! #        .unwrap();
//! #
//!     let array = py
//!         .eval(
//!             "np.array([np.datetime64('2017-04-21')])",
//!             None,
//!             Some(locals),
//!         )
//!         .unwrap()
//!         .downcast::<PyArray1<Datetime<units::Days>>>()
//!         .unwrap();
//!
//!     assert_eq!(
//!         array.get_owned(0).unwrap(),
//!         Datetime::<units::Days>::from(17_277)
//!     );
//!
//!     let array = py
//!         .eval(
//!             "np.array([np.datetime64('2022-03-29')]) - np.array([np.datetime64('2017-04-21')])",
//!             None,
//!             Some(locals),
//!         )
//!         .unwrap()
//!         .downcast::<PyArray1<Timedelta<units::Days>>>()
//!         .unwrap();
//!
//!     assert_eq!(
//!         array.get_owned(0).unwrap(),
//!         Timedelta::<units::Days>::from(1_803)
//!     );
//! });
//! ```
//!
//! [datetime]: https://numpy.org/doc/stable/reference/arrays.datetime.html
//! [scalars-datetime64]: https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.datetime64
//! [scalars-timedelta64]: https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.timedelta64

use std::cell::UnsafeCell;
use std::collections::hash_map::Entry;
use std::fmt;
use std::hash::Hash;
use std::marker::PhantomData;

use ahash::AHashMap;
use pyo3::{Py, Python};

use crate::dtype::{Element, PyArrayDescr};
use crate::npyffi::{PyArray_DatetimeDTypeMetaData, NPY_DATETIMEUNIT, NPY_TYPES};

/// Represents the [datetime units][datetime-units] supported by NumPy
///
/// [datetime-units]: https://numpy.org/doc/stable/reference/arrays.datetime.html#datetime-units
pub trait Unit: Send + Sync + Clone + Copy + PartialEq + Eq + Hash + PartialOrd + Ord {
    /// The matching NumPy [datetime unit code][NPY_DATETIMEUNIT]
    ///
    /// [NPY_DATETIMEUNIT]: https://github.com/numpy/numpy/blob/4c60b3263ac50e5e72f6a909e156314fc3c9cba0/numpy/core/include/numpy/ndarraytypes.h#L276
    const UNIT: NPY_DATETIMEUNIT;

    /// The abbrevation used for debug formatting
    const ABBREV: &'static str;
}

macro_rules! define_units {
    ($($(#[$meta:meta])* $struct:ident => $unit:ident $abbrev:literal,)+) => {
        $(

        $(#[$meta])*
        #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
        pub struct $struct;

        impl Unit for $struct {
            const UNIT: NPY_DATETIMEUNIT = NPY_DATETIMEUNIT::$unit;

            const ABBREV: &'static str = $abbrev;
        }

        )+
    };
}

/// Predefined implementors of the [`Unit`] trait
pub mod units {
    use super::*;

    define_units!(
        #[doc = "Years, i.e. 12 months"]
        Years => NPY_FR_Y "a",
        #[doc = "Months, i.e. 30 days"]
        Months => NPY_FR_M "mo",
        #[doc = "Weeks, i.e. 7 days"]
        Weeks => NPY_FR_W "w",
        #[doc = "Days, i.e. 24 hours"]
        Days => NPY_FR_D "d",
        #[doc = "Hours, i.e. 60 minutes"]
        Hours => NPY_FR_h "h",
        #[doc = "Minutes, i.e. 60 seconds"]
        Minutes => NPY_FR_m "min",
        #[doc = "Seconds"]
        Seconds => NPY_FR_s "s",
        #[doc = "Milliseconds, i.e. 10^-3 seconds"]
        Milliseconds => NPY_FR_ms "ms",
        #[doc = "Microseconds, i.e. 10^-6 seconds"]
        Microseconds => NPY_FR_us "µs",
        #[doc = "Nanoseconds, i.e. 10^-9 seconds"]
        Nanoseconds => NPY_FR_ns "ns",
        #[doc = "Picoseconds, i.e. 10^-12 seconds"]
        Picoseconds => NPY_FR_ps "ps",
        #[doc = "Femtoseconds, i.e. 10^-15 seconds"]
        Femtoseconds => NPY_FR_fs "fs",
        #[doc = "Attoseconds, i.e. 10^-18 seconds"]
        Attoseconds => NPY_FR_as "as",
    );
}

/// Corresponds to the [`datetime64`][scalars-datetime64] scalar type
///
/// [scalars-datetime64]: https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.datetime64
#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(transparent)]
pub struct Datetime(i64, PhantomData);

impl From<i64> for Datetime {
    fn from(val: i64) -> Self {
        Self(val, PhantomData)
    }
}

impl From<Datetime> for i64 {
    fn from(val: Datetime) -> Self {
        val.0
    }
}

unsafe impl Element for Datetime {
    const IS_COPY: bool = true;

    fn get_dtype(py: Python) -> &PyArrayDescr {
        static DTYPES: TypeDescriptors = unsafe { TypeDescriptors::new(NPY_TYPES::NPY_DATETIME) };

        DTYPES.from_unit(py, U::UNIT)
    }
}

impl fmt::Debug for Datetime {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Datetime({} {})", self.0, U::ABBREV)
    }
}

/// Corresponds to the [`timedelta64`][scalars-datetime64] scalar type
///
/// [scalars-timedelta64]: https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.timedelta64
#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(transparent)]
pub struct Timedelta(i64, PhantomData);

impl From<i64> for Timedelta {
    fn from(val: i64) -> Self {
        Self(val, PhantomData)
    }
}

impl From<Timedelta> for i64 {
    fn from(val: Timedelta) -> Self {
        val.0
    }
}

unsafe impl Element for Timedelta {
    const IS_COPY: bool = true;

    fn get_dtype(py: Python) -> &PyArrayDescr {
        static DTYPES: TypeDescriptors = unsafe { TypeDescriptors::new(NPY_TYPES::NPY_TIMEDELTA) };

        DTYPES.from_unit(py, U::UNIT)
    }
}

impl fmt::Debug for Timedelta {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Timedelta({} {})", self.0, U::ABBREV)
    }
}

struct TypeDescriptors {
    npy_type: NPY_TYPES,
    dtypes: UnsafeCell<Option<AHashMap<NPY_DATETIMEUNIT, Py<PyArrayDescr>>>>,
}

unsafe impl Sync for TypeDescriptors {}

impl TypeDescriptors {
    /// `npy_type` must be either `NPY_DATETIME` or `NPY_TIMEDELTA`.
    const unsafe fn new(npy_type: NPY_TYPES) -> Self {
        Self {
            npy_type,
            dtypes: UnsafeCell::new(None),
        }
    }

    #[allow(clippy::mut_from_ref)]
    unsafe fn get(&self) -> &mut AHashMap<NPY_DATETIMEUNIT, Py<PyArrayDescr>> {
        (*self.dtypes.get()).get_or_insert_with(AHashMap::new)
    }

    #[allow(clippy::wrong_self_convention)]
    fn from_unit<'py>(&'py self, py: Python<'py>, unit: NPY_DATETIMEUNIT) -> &'py PyArrayDescr {
        // SAFETY: We hold the GIL and we do not call into user code which might re-enter.
        let dtypes = unsafe { self.get() };

        match dtypes.entry(unit) {
            Entry::Occupied(entry) => entry.into_mut().as_ref(py),
            Entry::Vacant(entry) => {
                let dtype = PyArrayDescr::new_from_npy_type(py, self.npy_type);

                // SAFETY: `self.npy_type` is either `NPY_DATETIME` or `NPY_TIMEDELTA` which implies the type of `c_metadata`.
                unsafe {
                    let metadata = &mut *((*dtype.as_dtype_ptr()).c_metadata
                        as *mut PyArray_DatetimeDTypeMetaData);

                    metadata.meta.base = unit;
                    metadata.meta.num = 1;
                }

                entry.insert(dtype.into()).as_ref(py)
            }
        }
    }
}

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

    use pyo3::{
        py_run,
        types::{PyDict, PyModule},
    };

    use crate::array::PyArray1;

    #[test]
    fn from_python_to_rust() {
        Python::with_gil(|py| {
            let locals = py
                .eval("{ 'np': __import__('numpy') }", None, None)
                .unwrap()
                .downcast::<PyDict>()
                .unwrap();

            let array = py
                .eval(
                    "np.array([np.datetime64('1970-01-01')])",
                    None,
                    Some(locals),
                )
                .unwrap()
                .downcast::<PyArray1<Datetime<units::Days>>>()
                .unwrap();

            let value: i64 = array.get_owned(0).unwrap().into();
            assert_eq!(value, 0);
        });
    }

    #[test]
    fn from_rust_to_python() {
        Python::with_gil(|py| {
            let array = PyArray1::<Timedelta<units::Minutes>>::zeros(py, 1, false);

            *array.readwrite().get_mut(0).unwrap() = Timedelta::<units::Minutes>::from(5);

            let np = py
                .eval("__import__('numpy')", None, None)
                .unwrap()
                .downcast::<PyModule>()
                .unwrap();

            py_run!(py, array np, "assert array.dtype == np.dtype('timedelta64[m]')");
            py_run!(py, array np, "assert array[0] == np.timedelta64(5, 'm')");
        });
    }

    #[test]
    fn debug_formatting() {
        assert_eq!(
            format!("{:?}", Datetime::<units::Days>::from(28)),
            "Datetime(28 d)"
        );

        assert_eq!(
            format!("{:?}", Timedelta::<units::Milliseconds>::from(160)),
            "Timedelta(160 ms)"
        );
    }

    #[test]
    fn unit_conversion() {
        #[track_caller]
        fn convert<S: Unit, D: Unit>(py: Python<'_>, expected_value: i64) {
            let array = PyArray1::<Timedelta<S>>::from_slice(py, &[Timedelta::<S>::from(1)]);
            let array = array.cast::<Timedelta<D>>(false).unwrap();

            let value: i64 = array.get_owned(0).unwrap().into();
            assert_eq!(value, expected_value);
        }

        Python::with_gil(|py| {
            convert::<units::Years, units::Days>(py, (97 + 400 * 365) / 400);
            convert::<units::Months, units::Days>(py, (97 + 400 * 365) / 400 / 12);

            convert::<units::Weeks, units::Seconds>(py, 7 * 24 * 60 * 60);
            convert::<units::Days, units::Seconds>(py, 24 * 60 * 60);
            convert::<units::Hours, units::Seconds>(py, 60 * 60);
            convert::<units::Minutes, units::Seconds>(py, 60);

            convert::<units::Seconds, units::Milliseconds>(py, 1_000);
            convert::<units::Seconds, units::Microseconds>(py, 1_000_000);
            convert::<units::Seconds, units::Nanoseconds>(py, 1_000_000_000);
            convert::<units::Seconds, units::Picoseconds>(py, 1_000_000_000_000);
            convert::<units::Seconds, units::Femtoseconds>(py, 1_000_000_000_000_000);

            convert::<units::Femtoseconds, units::Attoseconds>(py, 1_000);
        });
    }
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::rc::Rc;
use std::cell::RefCell;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::error::Error;

use dbus::{Connection, NameFlag};
use dbus::tree::{Factory, Tree, Property, MethodFn, MethodErr, EmitsChangedSignal, Interface};
use dbus::MessageItem;
use dbus;

use blockdev::{BlockMember, BlockDevs};

use consts::*;

use stratis::Stratis;

use types::{StratisResult, StratisError};

#[derive(Debug, Clone)]
pub struct DbusContext<'a> {
    name_prop: Arc<Property<MethodFn<'a>>>,
    pub remaining_prop: Arc<Property<MethodFn<'a>>>,
    pub total_prop: Arc<Property<MethodFn<'a>>>,
    pub status_prop: Arc<Property<MethodFn<'a>>>,
    pub running_status_prop: Arc<Property<MethodFn<'a>>>,
    pub block_devices_prop: Arc<Property<MethodFn<'a>>>,
}

impl<'a> DbusContext<'a> {
    pub fn update_one(prop: &Arc<Property<MethodFn<'a>>>, m: MessageItem) -> StratisResult<()> {
        match prop.set_value(m) {
            Ok(_) => Ok(()), // TODO: return signals
            Err(()) => Err(StratisError::Dbus(dbus::Error::new_custom(
                "UpdateError", "Could not update property with value"))),
        }
    }
}

pub fn get_base_tree<'a>(c: &'a Connection) -> StratisResult<Tree<MethodFn<'a>>> {
    c.register_name(STRATIS_BASE_SERVICE, NameFlag::ReplaceExisting as u32).unwrap();

    let f = Factory::new_fn();

    let base_tree = f.tree();

    let listpools_method = f.method(LIST_POOLS, move |m, _, _| {
            let mut items = m.get_items();
            if items.len() < 1 {
                return Err(MethodErr::no_arg());
            }

            let name = try!(items.pop()
                .ok_or_else(MethodErr::no_arg)
                .and_then(|i| {
                    i.inner::<&str>()
                        .map_err(|_| MethodErr::invalid_arg(&i))
                        .map(|i| i.to_owned())
                }));

            println!("method called");

            Ok(vec![m.method_return()])
        })
        .in_arg(("name", "s"));
    let createpool_method = f.method(CREATE_POOL, move |m, _, _| {
            let mut items = m.get_items();
            if items.len() < 1 {
                return Err(MethodErr::no_arg());
            }

            let name = try!(items.pop()
                .ok_or_else(MethodErr::no_arg)
                .and_then(|i| {
                    i.inner::<&str>()
                        .map_err(|_| MethodErr::invalid_arg(&i))
                        .map(|i| i.to_owned())
                }));

            println!("method called");

            Ok(vec![m.method_return()])
        })
        .in_arg(("name", "s"));
    let destroypool_method = f.method(DESTROY_POOL, move |m, _, _| {
            let mut items = m.get_items();
            if items.len() < 1 {
                return Err(MethodErr::no_arg());
            }

            let name = try!(items.pop()
                .ok_or_else(MethodErr::no_arg)
                .and_then(|i| {
                    i.inner::<&str>()
                        .map_err(|_| MethodErr::invalid_arg(&i))
                        .map(|i| i.to_owned())
                }));

            println!("method called");

            Ok(vec![m.method_return()])
        })
        .in_arg(("name", "s"));
    let getpoolobjectpath_method = f.method(GET_POOL_OBJECT_PATH, move |m, _, _| {
            let mut items = m.get_items();
            if items.len() < 1 {
                return Err(MethodErr::no_arg());
            }

            let name = try!(items.pop()
                .ok_or_else(MethodErr::no_arg)
                .and_then(|i| {
                    i.inner::<&str>()
                        .map_err(|_| MethodErr::invalid_arg(&i))
                        .map(|i| i.to_owned())
                }));

            println!("method called");

            Ok(vec![m.method_return()])
        })
        .in_arg(("name", "s"));
    let getvolumeobjectpath_method = f.method(GET_VOLUME_OBJECT_PATH, move |m, _, _| {
            let mut items = m.get_items();
            if items.len() < 1 {
                return Err(MethodErr::no_arg());
            }

            let name = try!(items.pop()
                .ok_or_else(MethodErr::no_arg)
                .and_then(|i| {
                    i.inner::<&str>()
                        .map_err(|_| MethodErr::invalid_arg(&i))
                        .map(|i| i.to_owned())
                }));

            println!("method called");

            Ok(vec![m.method_return()])
        })
        .in_arg(("name", "s"));
    let getdevobjectpath_method = f.method(GET_DEV_OBJECT_PATH, move |m, _, _| {
            let mut items = m.get_items();
            if items.len() < 1 {
                return Err(MethodErr::no_arg());
            }

            let name = try!(items.pop()
                .ok_or_else(MethodErr::no_arg)
                .and_then(|i| {
                    i.inner::<&str>()
                        .map_err(|_| MethodErr::invalid_arg(&i))
                        .map(|i| i.to_owned())
                }));

            println!("method called");

            Ok(vec![m.method_return()])
        })
        .in_arg(("name", "s"));
    let getcacheobjectpath_method = f.method(GET_CACHE_OBJECT_PATH, move |m, _, _| {
            let mut items = m.get_items();
            if items.len() < 1 {
                return Err(MethodErr::no_arg());
            }

            let name = try!(items.pop()
                .ok_or_else(MethodErr::no_arg)
                .and_then(|i| {
                    i.inner::<&str>()
                        .map_err(|_| MethodErr::invalid_arg(&i))
                        .map(|i| i.to_owned())
                }));

            println!("method called");

            Ok(vec![m.method_return()])
        })
        .in_arg(("name", "s"));
    let geterrorcodes_method = f.method(GET_ERROR_CODES, move |m, _, _| {
            let mut items = m.get_items();
            if items.len() < 1 {
                return Err(MethodErr::no_arg());
            }

            let name = try!(items.pop()
                .ok_or_else(MethodErr::no_arg)
                .and_then(|i| {
                    i.inner::<&str>()
                        .map_err(|_| MethodErr::invalid_arg(&i))
                        .map(|i| i.to_owned())
                }));

            println!("method called");

            Ok(vec![m.method_return()])
        })
        .in_arg(("name", "s"));
    let getraidlevels_method = f.method(GET_RAID_LEVELS, move |m, _, _| {
            let mut items = m.get_items();
            if items.len() < 1 {
                return Err(MethodErr::no_arg());
            }

            let name = try!(items.pop()
                .ok_or_else(MethodErr::no_arg)
                .and_then(|i| {
                    i.inner::<&str>()
                        .map_err(|_| MethodErr::invalid_arg(&i))
                        .map(|i| i.to_owned())
                }));

            println!("method called");

            Ok(vec![m.method_return()])
        })
        .in_arg(("name", "s"));
    let getdevtypes_method = f.method(GET_DEV_TYPES, move |m, _, _| {
            let mut items = m.get_items();
            if items.len() < 1 {
                return Err(MethodErr::no_arg());
            }

            let name = try!(items.pop()
                .ok_or_else(MethodErr::no_arg)
                .and_then(|i| {
                    i.inner::<&str>()
                        .map_err(|_| MethodErr::invalid_arg(&i))
                        .map(|i| i.to_owned())
                }));

            println!("method called");

            Ok(vec![m.method_return()])
        })
        .in_arg(("name", "s"));

    let obj_path = f.object_path(STRATIS_BASE_PATH)
        .introspectable()
        .add(f.interface(STRATIS_BASE_SERVICE)
        	.add_m(listpools_method)
			.add_m(createpool_method)
			.add_m(destroypool_method)
			.add_m(getpoolobjectpath_method)
			.add_m(getvolumeobjectpath_method)
			.add_m(getdevobjectpath_method)
			.add_m(getcacheobjectpath_method)
			.add_m(geterrorcodes_method)
			.add_m(getraidlevels_method)
			.add_m(getdevtypes_method));

    let base_tree = base_tree.add(obj_path);
    try!(base_tree.set_registered(c, true));

    Ok(base_tree)
}
<commit_msg>split methods out so they are not constructed inline<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::rc::Rc;
use std::cell::RefCell;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::error::Error;

use dbus::{Connection, NameFlag};
use dbus::tree::{Factory, Tree, Property, MethodFn, MethodErr, EmitsChangedSignal, Interface};
use dbus::MessageItem;
use dbus;
use dbus::Message;
use dbus::tree::MethodResult;

use blockdev::{BlockMember, BlockDevs};

use consts::*;

use stratis::Stratis;

use types::{StratisResult, StratisError};

#[derive(Debug, Clone)]
pub struct DbusContext<'a> {
    name_prop: Arc<Property<MethodFn<'a>>>,
    pub remaining_prop: Arc<Property<MethodFn<'a>>>,
    pub total_prop: Arc<Property<MethodFn<'a>>>,
    pub status_prop: Arc<Property<MethodFn<'a>>>,
    pub running_status_prop: Arc<Property<MethodFn<'a>>>,
    pub block_devices_prop: Arc<Property<MethodFn<'a>>>,
}

impl<'a> DbusContext<'a> {
    pub fn update_one(prop: &Arc<Property<MethodFn<'a>>>, m: MessageItem) -> StratisResult<()> {
        match prop.set_value(m) {
            Ok(_) => Ok(()), // TODO: return signals
            Err(()) => Err(StratisError::Dbus(dbus::Error::new_custom(
                "UpdateError", "Could not update property with value"))),
        }
    }
}

fn listpools(m: &Message) -> MethodResult {
    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    println!("method called");

    Ok(vec![m.method_return()])
}

fn createpool(m: &Message) -> MethodResult {
    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    println!("method called");

    Ok(vec![m.method_return()])
}

fn destroypool(m: &Message) -> MethodResult {
    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    println!("method called");

    Ok(vec![m.method_return()])
}

fn getpoolobjectpath(m: &Message) -> MethodResult {
    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    println!("method called");

    Ok(vec![m.method_return()])
}

fn getvolumeobjectpath(m: &Message) -> MethodResult {
    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    println!("method called");

    Ok(vec![m.method_return()])
}

fn getdevobjectpath(m: &Message) -> MethodResult {
    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    println!("method called");

    Ok(vec![m.method_return()])
}

fn getcacheobjectpath(m: &Message) -> MethodResult {
    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    println!("method called");

    Ok(vec![m.method_return()])
}

fn geterrorcodes(m: &Message) -> MethodResult {
    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    println!("method called");

    Ok(vec![m.method_return()])
}

fn getraidlevels(m: &Message) -> MethodResult {
    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    println!("method called");

    Ok(vec![m.method_return()])
}

fn getdevtypes(m: &Message) -> MethodResult {
    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    println!("method called");

    Ok(vec![m.method_return()])
}
pub fn get_base_tree<'a>(c: &'a Connection) -> StratisResult<Tree<MethodFn<'a>>> {
    c.register_name(STRATIS_BASE_SERVICE, NameFlag::ReplaceExisting as u32).unwrap();

    let f = Factory::new_fn();

    let base_tree = f.tree();

    let listpools_method = f.method(LIST_POOLS, move |m, _, _| { listpools(m) })
         .in_arg(("poolname", "s"))
         .in_arg(("devlist", "as"))
         .in_arg(("raidtype", "i"))
         .out_arg(("object_name", "s"))
         .out_arg(("return_code", "i"))
         .out_arg(("return_string", "s"));

    let createpool_method = f.method(CREATE_POOL, move |m, _, _| { createpool(m) })
         .in_arg(("name", "s"));

    let destroypool_method = f.method(DESTROY_POOL, move |m, _, _| { destroypool(m) })
         .in_arg(("name", "s"));

    let getpoolobjectpath_method = f.method(GET_POOL_OBJECT_PATH, move |m, _, _| { getpoolobjectpath(m) })
         .in_arg(("name", "s"));

    let getvolumeobjectpath_method = f.method(GET_VOLUME_OBJECT_PATH, move |m, _, _| { getvolumeobjectpath(m) })
         .in_arg(("name", "s"));

    let getdevobjectpath_method = f.method(GET_DEV_OBJECT_PATH, move |m, _, _| { getdevobjectpath(m) })
         .in_arg(("name", "s"));

    let getcacheobjectpath_method = f.method(GET_CACHE_OBJECT_PATH, move |m, _, _| { getcacheobjectpath(m) })
         .in_arg(("name", "s"));

    let geterrorcodes_method = f.method(GET_ERROR_CODES, move |m, _, _| { geterrorcodes(m) })
         .in_arg(("name", "s"));

    let getraidlevels_method = f.method(GET_RAID_LEVELS, move |m, _, _| { getraidlevels(m) })
         .in_arg(("name", "s"));

    let getdevtypes_method = f.method(GET_DEV_TYPES, move |m, _, _| { getdevtypes(m) })
         .in_arg(("name", "s"));

    let obj_path = f.object_path(STRATIS_BASE_PATH)
        .introspectable()
        .add(f.interface(STRATIS_BASE_SERVICE)
            .add_m(listpools_method)
            .add_m(createpool_method)
            .add_m(destroypool_method)
            .add_m(getpoolobjectpath_method)
            .add_m(getvolumeobjectpath_method)
            .add_m(getdevobjectpath_method)
            .add_m(getcacheobjectpath_method)
            .add_m(geterrorcodes_method)
            .add_m(getraidlevels_method)
            .add_m(getdevtypes_method));

    let base_tree = base_tree.add(obj_path);
    try!(base_tree.set_registered(c, true));

    Ok(base_tree)
}
<|endoftext|>


<commit_before>use clap::ArgMatches;

use config::Config;
use libdo::DoManager;

pub fn run(m: &ArgMatches, cfg: &Config) {
    match m.subcommand() {
        ("list-actions", _) => {
            // PLACEHOLDER
            println!("Listing all account actions");
        },
        ("action", Some(m)) => {
            // PLACEHOLDER
            println!("Retrieving action id: {}", m.value_of("id").unwrap());
        },
        ("", None)               => {
            if cfg.debug { println!("Displaying account with token:\n\t{}\n", &cfg.auth[..])}
            let domgr = DoManager::with_token(&cfg.auth[..]);
            if cfg.debug { println!("Sending request:\n\t{}\n", domgr.account.to_string().replace("\n", "\n\t")); }
            if cfg.no_send { return }
            match domgr.account.retrieve() {
                Ok(s) => println!("{}", s),
                Err(e) => println!("{}", e)
            }
        },
        _ => unreachable!()
    }
}
<commit_msg>feat(account): implements account subcommand<commit_after>use clap::ArgMatches;

use config::Config;
use libdo::DoManager;

pub fn run(m: &ArgMatches, cfg: &Config) {
    let domgr = DoManager::with_token(&cfg.auth[..]);
    if cfg.debug { println!("Displaying account with token:\n\t{}\n", &cfg.auth[..]) }
    match m.subcommand() {
        ("list-actions", _) => {
            if cfg.debug { println!("Displaying all account actions\n") }
            if cfg.debug {
                println!("Sending request:\n\t{}\n", domgr.account.actions()
                                                                  .to_string()
                                                                  .replace("\n", "\n\t"));
            }
            if cfg.no_send { return }
            match domgr.account.retrieve_actions() {
                Ok(v) => {
                    for act in v {
                        println!("{}", act)
                    }
                },
                Err(e) => println!("{}", e)
            }
        },
        ("action", Some(m)) => {
            let id = m.value_of("id").unwrap();
            if cfg.debug { println!("Displaying action with ID: {}\n", id) }
            if cfg.debug {
                println!("Sending request:\n\t{}\n", domgr.account.action(id)
                                                                  .to_string()
                                                                  .replace("\n", "\n\t"));
            }
            if cfg.no_send { return }
            match domgr.account.retrieve_action(id) {
                Ok(s) => println!("{}", s),
                Err(e) => println!("{}", e)
            }
        },
        ("", None)               => {
            if cfg.debug { println!("Sending request:\n\t{}\n", domgr.account.to_string().replace("\n", "\n\t")); }
            if cfg.no_send { return }
            match domgr.account.retrieve() {
                Ok(s) => println!("{}", s),
                Err(e) => println!("{}", e)
            }
        },
        _ => unreachable!()
    }
}
<|endoftext|>


<commit_before>#[macro_escape];

use std::io::BufferedWriter;
use std::io::{File, Writer};
use std::str::Chars;
use std::vec_ng::Vec;

struct ParseBranch {
    matches: Vec<u8>,
    result: Option<~str>,
    children: Vec<ParseBranch>,
}

impl ParseBranch {
    fn new() -> ParseBranch {
        ParseBranch {
            matches: Vec::new(),
            result: None,
            children: Vec::new(),
        }
    }
}

pub fn branchify(options: &[(&str, &str)], case_sensitive: bool) -> Vec<ParseBranch> {
    let mut root = ParseBranch::new();

    fn go_down_moses(branch: &mut ParseBranch, mut chariter: Chars, result: &str, case_sensitive: bool) {
        match chariter.next() {
            Some(c) => {
                let first_case = if case_sensitive { c as u8 } else { c.to_ascii().to_upper().to_byte() };
                for next_branch in branch.children.mut_iter() {
                    if next_branch.matches.as_slice()[0] == first_case {
                        go_down_moses(next_branch, chariter, result, case_sensitive);
                        return;
                    }
                }
                let mut subbranch = ParseBranch::new();
                subbranch.matches.push(first_case);
                if !case_sensitive {
                    let second_case = c.to_ascii().to_lower().to_byte();
                    if first_case != second_case {
                        subbranch.matches.push(second_case);
                    }
                }
                branch.children.push(subbranch);
                let index = branch.children.len() -1;
                go_down_moses(&mut branch.children.as_mut_slice()[index], chariter, result, case_sensitive);
            },
            None => {
                assert!(branch.result.is_none());
                branch.result = Some(result.to_owned());
            },
        }
    };

    for &(key, result) in options.iter() {
        go_down_moses(&mut root, key.chars(), result, case_sensitive);
    }

    root.children
}

macro_rules! branchify(
    (case sensitive, $($key:expr => $value:ident),*) => (
        ::branchify::branchify([$(($key, stringify!($value))),*], true)
    );
    (case insensitive, $($key:expr => $value:ident),*) => (
        branchify([$(($key, stringify!($value))),*], false)
    );
)

/// Prints the contents to stdout.
///
/// :param branches: the branches to search through
/// :param indent: the level of indentation (each level representing four leading spaces)
/// :param read_call: the function call to read a byte
/// :param end: the byte which marks the end of the sequence
/// :param max_len: the maximum length a value may be before giving up and returning ``None``
/// :param valid: the function call to if a byte ``b`` is valid
/// :param unknown: the expression to call for an unknown value; in this string, ``{}`` will be
///         replaced with an expression (literal or non-literal) evaluating to a ``~str`` (it is
///         ``{}`` only, not arbitrary format strings)
pub fn generate_branchified_method(
        writer: &mut BufferedWriter<File>,
        branches: &[ParseBranch],
        indent: uint,
        read_call: &str,
        end: &str,
        max_len: &str,
        valid: &str,
        unknown: &str) {

    // Write Formatted
    macro_rules! wf(($($x:tt)*) => ({
        let indentstr = " ".repeat(indent * 4);
        let s = format!($($x)*);
        let result = writer.write(indentstr.as_bytes())
        .and(writer.write(s.as_bytes()))
        .and(writer.write(bytes!("\n")));

        match result {
            Ok(_)  => {},
            Err(e) => fail!("write error: {:s}", e.desc),
        }
    }))

    fn r(writer: &mut BufferedWriter<File>, branch: &ParseBranch, prefix: &str, indent: uint, read_call: &str,
            end: &str, max_len: &str, valid: &str, unknown: &str) {
        for &c in branch.matches.iter() {
            let next_prefix = format!("{}{}", prefix, c as char);
            wf!("Some(b) if b == '{}' as u8 => match {} \\{", c as char, read_call);
            for b in branch.children.iter() {
                r(writer, b, next_prefix, indent + 1, read_call, end, max_len, valid, unknown);
            }
            match branch.result {
                Some(ref result) => wf!("    Some(b) if b == SP => return Some({}),", *result),
                None => wf!("    Some(b) if b == SP => return Some({}),",
                                unknown.replace("{}", format!("~\"{}\"", next_prefix))),
            }
            wf!("    Some(b) if {} => (\"{}\", b),", valid, next_prefix);
            wf!("    _ => return None,");
            wf!("\\},");
        }
    }
    wf!("let (s, next_byte) = match {} \\{", read_call);
    for b in branches.iter() {
        r(writer, b, "", indent + 1, read_call, end, max_len, valid, unknown);
    }
    wf!("    Some(b) if {} => (\"\", b),", valid);
    wf!("    _ => return None,");
    wf!("\\};");
    wf!("// OK, that didn't pan out. Let's read the rest and see what we get.");
    wf!("let mut s = s.to_owned();");
    wf!("s.push_char(next_byte as char);");
    wf!("loop \\{");
    wf!("    match {} \\{", read_call);
    wf!("        Some(b) if b == {} => return Some({}),", end, unknown.replace("{}", "s"));
    wf!("        Some(b) if {} => \\{", valid);
    wf!("            if s.len() == {} \\{", max_len);
    wf!("                // Too long; bad request");
    wf!("                return None;");
    wf!("            \\}");
    wf!("            s.push_char(b as char);");
    wf!("        \\},");
    wf!("        _ => return None,");
    wf!("    \\}");
    wf!("\\}");
}
<commit_msg>Suppress warnings about visible private types<commit_after>#[macro_escape];

use std::io::BufferedWriter;
use std::io::{File, Writer};
use std::str::Chars;
use std::vec_ng::Vec;

struct ParseBranch {
    matches: Vec<u8>,
    result: Option<~str>,
    children: Vec<ParseBranch>,
}

impl ParseBranch {
    fn new() -> ParseBranch {
        ParseBranch {
            matches: Vec::new(),
            result: None,
            children: Vec::new(),
        }
    }
}

#[allow(visible_private_types)]
pub fn branchify(options: &[(&str, &str)], case_sensitive: bool) -> Vec<ParseBranch> {
    let mut root = ParseBranch::new();

    fn go_down_moses(branch: &mut ParseBranch, mut chariter: Chars, result: &str, case_sensitive: bool) {
        match chariter.next() {
            Some(c) => {
                let first_case = if case_sensitive { c as u8 } else { c.to_ascii().to_upper().to_byte() };
                for next_branch in branch.children.mut_iter() {
                    if next_branch.matches.as_slice()[0] == first_case {
                        go_down_moses(next_branch, chariter, result, case_sensitive);
                        return;
                    }
                }
                let mut subbranch = ParseBranch::new();
                subbranch.matches.push(first_case);
                if !case_sensitive {
                    let second_case = c.to_ascii().to_lower().to_byte();
                    if first_case != second_case {
                        subbranch.matches.push(second_case);
                    }
                }
                branch.children.push(subbranch);
                let index = branch.children.len() -1;
                go_down_moses(&mut branch.children.as_mut_slice()[index], chariter, result, case_sensitive);
            },
            None => {
                assert!(branch.result.is_none());
                branch.result = Some(result.to_owned());
            },
        }
    };

    for &(key, result) in options.iter() {
        go_down_moses(&mut root, key.chars(), result, case_sensitive);
    }

    root.children
}

macro_rules! branchify(
    (case sensitive, $($key:expr => $value:ident),*) => (
        ::branchify::branchify([$(($key, stringify!($value))),*], true)
    );
    (case insensitive, $($key:expr => $value:ident),*) => (
        branchify([$(($key, stringify!($value))),*], false)
    );
)

/// Prints the contents to stdout.
///
/// :param branches: the branches to search through
/// :param indent: the level of indentation (each level representing four leading spaces)
/// :param read_call: the function call to read a byte
/// :param end: the byte which marks the end of the sequence
/// :param max_len: the maximum length a value may be before giving up and returning ``None``
/// :param valid: the function call to if a byte ``b`` is valid
/// :param unknown: the expression to call for an unknown value; in this string, ``{}`` will be
///         replaced with an expression (literal or non-literal) evaluating to a ``~str`` (it is
///         ``{}`` only, not arbitrary format strings)
#[allow(visible_private_types)]
pub fn generate_branchified_method(
        writer: &mut BufferedWriter<File>,
        branches: &[ParseBranch],
        indent: uint,
        read_call: &str,
        end: &str,
        max_len: &str,
        valid: &str,
        unknown: &str) {

    // Write Formatted
    macro_rules! wf(($($x:tt)*) => ({
        let indentstr = " ".repeat(indent * 4);
        let s = format!($($x)*);
        let result = writer.write(indentstr.as_bytes())
        .and(writer.write(s.as_bytes()))
        .and(writer.write(bytes!("\n")));

        match result {
            Ok(_)  => {},
            Err(e) => fail!("write error: {:s}", e.desc),
        }
    }))

    fn r(writer: &mut BufferedWriter<File>, branch: &ParseBranch, prefix: &str, indent: uint, read_call: &str,
            end: &str, max_len: &str, valid: &str, unknown: &str) {
        for &c in branch.matches.iter() {
            let next_prefix = format!("{}{}", prefix, c as char);
            wf!("Some(b) if b == '{}' as u8 => match {} \\{", c as char, read_call);
            for b in branch.children.iter() {
                r(writer, b, next_prefix, indent + 1, read_call, end, max_len, valid, unknown);
            }
            match branch.result {
                Some(ref result) => wf!("    Some(b) if b == SP => return Some({}),", *result),
                None => wf!("    Some(b) if b == SP => return Some({}),",
                                unknown.replace("{}", format!("~\"{}\"", next_prefix))),
            }
            wf!("    Some(b) if {} => (\"{}\", b),", valid, next_prefix);
            wf!("    _ => return None,");
            wf!("\\},");
        }
    }
    wf!("let (s, next_byte) = match {} \\{", read_call);
    for b in branches.iter() {
        r(writer, b, "", indent + 1, read_call, end, max_len, valid, unknown);
    }
    wf!("    Some(b) if {} => (\"\", b),", valid);
    wf!("    _ => return None,");
    wf!("\\};");
    wf!("// OK, that didn't pan out. Let's read the rest and see what we get.");
    wf!("let mut s = s.to_owned();");
    wf!("s.push_char(next_byte as char);");
    wf!("loop \\{");
    wf!("    match {} \\{", read_call);
    wf!("        Some(b) if b == {} => return Some({}),", end, unknown.replace("{}", "s"));
    wf!("        Some(b) if {} => \\{", valid);
    wf!("            if s.len() == {} \\{", max_len);
    wf!("                // Too long; bad request");
    wf!("                return None;");
    wf!("            \\}");
    wf!("            s.push_char(b as char);");
    wf!("        \\},");
    wf!("        _ => return None,");
    wf!("    \\}");
    wf!("\\}");
}
<|endoftext|>


<commit_before>use std::ascii::AsciiExt;
use std::sync;
use std::thread;

use interface;
use irc;
use client;
use events;
use fern;

pub struct Dispatch {
    interface: interface::IrcInterface,
    state: sync::Arc<client::Client>,
    data_in: Receiver<irc::IrcMessage>,
    workers_out: Sender<PluginTask>,
}

impl Dispatch {
    pub fn new(interface: interface::IrcInterface, state: sync::Arc<client::Client>, data_in: Receiver<irc::IrcMessage>, logger: fern::ArcLogger) -> Dispatch {
        let (dispatch_out, workers_in) = channel();
        let workers_in_arc = sync::Arc::new(sync::Mutex::new(workers_in));
        for _ in range::<u8>(0, 4) {
            let executor = PluginExecutor::new(interface.clone(), workers_in_arc.clone(), logger.clone());
            executor.start_worker_thread();
        }

        return Dispatch {
            interface: interface,
            state: state,
            data_in: data_in,
            workers_out: dispatch_out,
        };
    }

    pub fn dispatch_loop(self) {
        loop {
            let message = match self.data_in.recv_opt() {
                Ok(v) => v,
                Err(()) => break,
            };
            self.process_message(&message);
        }
    }

    pub fn start_dispatch_loop(self) -> thread::Result<()>{
        return thread::Builder::new().name("dispatch".to_string()).spawn(move || {
            self.dispatch_loop();
        }).join();
    }

    // Noting: This has to be a separate method from spawn_dispatch_thread, so that we can name an 'a lifetime.
    // This allows us to give the new &str slices a specific lifetime, which I don't know a way to do without making a new function.
    fn process_message<'a>(&self, message: &'a irc::IrcMessage) {
        let plugins = self.state.plugins.read().unwrap();

        // PING
        if message.command.as_slice().eq_ignore_ascii_case("PING") {
            self.interface.send_raw(format!("PONG {}", message.args.connect(" ")));
        }

        let message_event = events::MessageTransport::from_internal(message);

        // Catch all listeners
        for listener in plugins.catch_all.iter() {
            self.execute(PluginTask::Message((listener.clone(), message_event.clone())));
        }

        // Raw listeners
        if let Some(list) = plugins.raw_listeners.get(&message.command.to_ascii_lowercase()) {
            for listener in list.iter() {
                self.execute(PluginTask::Message((listener.clone(), message_event.clone())));
            }
        }

        if message.command.as_slice().eq_ignore_ascii_case("PRIVMSG") {
            let channel = message.channel.as_ref().unwrap().as_slice(); // Always exists for PRIVMSG

            // CTCP
            if let Some(ctcp_event) = events::CtcpTransport::from_internal(message) {
                if let Some(list) = plugins.ctcp_listeners.get(&message.args[0].to_ascii_lowercase()) {
                    for ctcp_listener in list.iter() {
                        self.execute(PluginTask::Ctcp((ctcp_listener.clone(), ctcp_event.clone())));
                    }
                }
            }

            // Commands
            let command_prefix = format!(":{}", self.state.command_prefix.as_slice());

            // This checks for the command prefix, commands typed like '.command_name args'
            if message.args[1].starts_with(command_prefix.as_slice()) {
                let command = message.args[1].slice_from(command_prefix.len());
                let args = message.args.slice_from(2).iter().map(|s| s.clone()).collect::<Vec<String>>();
                self.dispatch_command(&plugins, command, channel, args, &message.mask);
            } else {
                // This checks for someone typing commands like 'BotName, command_name args'
                // We store whether or not a command was matched in a variable so that we can use it below.
                let mut command_matched = false;
                if let Some(captures) = regex!(r"^:([^\s]+?)[:;,]?\s+(.+)$").captures(message.args.slice_from(1).connect(" ").as_slice()) {
                    let same = {
                        let state = self.state.state.read().unwrap();
                        captures.at(1) == Some(state.nick.as_slice())
                    };
                    if same {
                        if let Some(args_str) = captures.at(2) {
                            let split = args_str.split(' ').collect::<Vec<&str>>();
                            let command = split[0];
                            let args = split.slice_from(1).iter().map(|s| s.to_string()).collect::<Vec<String>>();
                            self.dispatch_command(&plugins, command, channel, args, &message.mask);
                            command_matched = true;
                        }
                    }
                }

                // This checks for commands in a private message, where a prefix isn't required
                // People can just say 'command args' in a private message.
                // If the channel is the sender's nick, the message is being sent in a private message.
                if !command_matched && message.mask.nick() == Some(channel) {
                    let command = message.args[1].slice_from(1); // slice_from(1) to remove the `:` at the beginning of privmsg content.
                    let args = message.args.slice_from(2).iter().map(|s| s.clone()).collect::<Vec<String>>();
                    self.dispatch_command(&plugins, command, channel, args, &message.mask);
                }
            }
        }
    }

    fn dispatch_command(&self, plugins: &sync::RWLockReadGuard<client::PluginRegister>, command: &str, channel: &str, args: Vec<String>, mask: &irc::IrcMask) {
        if let Some(list) = plugins.commands.get(&command.to_ascii_lowercase()) {
            let command_event = events::CommandTransport::new(channel, args, mask);
            for closure in list.iter() {
                self.execute(PluginTask::Command((closure.clone(), command_event.clone())));
            }
        }
    }

    fn execute(&self, task: PluginTask) {
        self.workers_out.send(task);
    }
}

/// TODO: Better name for this
enum PluginTask {
    Command((sync::Arc<client::CommandListener>, events::CommandTransport)),
    Message((sync::Arc<client::MessageListener>, events::MessageTransport)),
    Ctcp((sync::Arc<client::CtcpListener>, events::CtcpTransport)),
}

impl PluginTask {
    fn execute(self, interface: &interface::IrcInterface) {
        match self {
            PluginTask::Command((closure, event)) => {
                closure.call((&events::CommandEvent::new(interface, &event),));
            },
            PluginTask::Message((closure, event)) => {
                closure.call((&events::MessageEvent::new(interface, &event),));
            },
            PluginTask::Ctcp((closure, event)) => {
                closure.call((&events::CtcpEvent::new(interface, &event),));
            },
        }
    }
}

struct PluginExecutor {
    interface: interface::IrcInterface,
    data_in: sync::Arc<sync::Mutex<Receiver<PluginTask>>>,
    logger: fern::ArcLogger,
    active: bool,
}

impl PluginExecutor {
    fn new(interface: interface::IrcInterface, data_in: sync::Arc<sync::Mutex<Receiver<PluginTask>>>, logger: fern::ArcLogger) -> PluginExecutor {
        return PluginExecutor {
            interface: interface,
            data_in: data_in,
            logger: logger,
            active: true,
        };
    }

    fn worker_loop(&mut self) {
        loop {
            let message = {
                // Only lock jobs for the time it takes
                // to get a job, not run it.
                let lock = self.data_in.lock().unwrap();
                lock.recv_opt()
            };
            match message {
                Ok(next) => next.execute(&self.interface),
                Err(()) => {
                    self.active = false;
                    break;
                }
            };
        }
    }

    fn start_worker_thread(mut self) {
        thread::Builder::new().name("worker_thread".to_string()).spawn(move || {
            fern::local::set_thread_logger(self.logger.clone());
            self.worker_loop();
        }).detach();
    }
}

impl Drop for PluginExecutor {
    fn drop(&mut self) {
        if self.active {
            warning!("Worker panicked!");
            PluginExecutor::new(self.interface.clone(), self.data_in.clone(), self.logger.clone()).start_worker_thread();
        }
    }
}
<commit_msg>Fix CTCP events<commit_after>use std::ascii::AsciiExt;
use std::sync;
use std::thread;

use interface;
use irc;
use client;
use events;
use fern;

pub struct Dispatch {
    interface: interface::IrcInterface,
    state: sync::Arc<client::Client>,
    data_in: Receiver<irc::IrcMessage>,
    workers_out: Sender<PluginTask>,
}

impl Dispatch {
    pub fn new(interface: interface::IrcInterface, state: sync::Arc<client::Client>, data_in: Receiver<irc::IrcMessage>, logger: fern::ArcLogger) -> Dispatch {
        let (dispatch_out, workers_in) = channel();
        let workers_in_arc = sync::Arc::new(sync::Mutex::new(workers_in));
        for _ in range::<u8>(0, 4) {
            let executor = PluginExecutor::new(interface.clone(), workers_in_arc.clone(), logger.clone());
            executor.start_worker_thread();
        }

        return Dispatch {
            interface: interface,
            state: state,
            data_in: data_in,
            workers_out: dispatch_out,
        };
    }

    pub fn dispatch_loop(self) {
        loop {
            let message = match self.data_in.recv_opt() {
                Ok(v) => v,
                Err(()) => break,
            };
            self.process_message(&message);
        }
    }

    pub fn start_dispatch_loop(self) -> thread::Result<()>{
        return thread::Builder::new().name("dispatch".to_string()).spawn(move || {
            self.dispatch_loop();
        }).join();
    }

    // Noting: This has to be a separate method from spawn_dispatch_thread, so that we can name an 'a lifetime.
    // This allows us to give the new &str slices a specific lifetime, which I don't know a way to do without making a new function.
    fn process_message<'a>(&self, message: &'a irc::IrcMessage) {
        let plugins = self.state.plugins.read().unwrap();

        // PING
        if message.command.as_slice().eq_ignore_ascii_case("PING") {
            self.interface.send_raw(format!("PONG {}", message.args.connect(" ")));
        }

        let message_event = events::MessageTransport::from_internal(message);

        // Catch all listeners
        for listener in plugins.catch_all.iter() {
            self.execute(PluginTask::Message((listener.clone(), message_event.clone())));
        }

        // Raw listeners
        if let Some(list) = plugins.raw_listeners.get(&message.command.to_ascii_lowercase()) {
            for listener in list.iter() {
                self.execute(PluginTask::Message((listener.clone(), message_event.clone())));
            }
        }

        if message.command.as_slice().eq_ignore_ascii_case("PRIVMSG") {
            let channel = message.channel.as_ref().unwrap().as_slice(); // Always exists for PRIVMSG

            // CTCP
            if let Some(ctcp_event) = events::CtcpTransport::from_internal(message) {
                if let Some(list) = plugins.ctcp_listeners.get(&ctcp_event.command
                        .to_ascii_lowercase()) {
                    for ctcp_listener in list.iter() {
                        self.execute(PluginTask::Ctcp((ctcp_listener.clone(),
                            ctcp_event.clone())));
                    }
                }
            }

            // Commands
            let command_prefix = format!(":{}", self.state.command_prefix.as_slice());

            // This checks for the command prefix, commands typed like '.command_name args'
            if message.args[1].starts_with(command_prefix.as_slice()) {
                let command = message.args[1].slice_from(command_prefix.len());
                let args = message.args.slice_from(2).iter().map(|s| s.clone()).collect::<Vec<String>>();
                self.dispatch_command(&plugins, command, channel, args, &message.mask);
            } else {
                // This checks for someone typing commands like 'BotName, command_name args'
                // We store whether or not a command was matched in a variable so that we can use it below.
                let mut command_matched = false;
                if let Some(captures) = regex!(r"^:([^\s]+?)[:;,]?\s+(.+)$").captures(message.args.slice_from(1).connect(" ").as_slice()) {
                    let same = {
                        let state = self.state.state.read().unwrap();
                        captures.at(1) == Some(state.nick.as_slice())
                    };
                    if same {
                        if let Some(args_str) = captures.at(2) {
                            let split = args_str.split(' ').collect::<Vec<&str>>();
                            let command = split[0];
                            let args = split.slice_from(1).iter().map(|s| s.to_string()).collect::<Vec<String>>();
                            self.dispatch_command(&plugins, command, channel, args, &message.mask);
                            command_matched = true;
                        }
                    }
                }

                // This checks for commands in a private message, where a prefix isn't required
                // People can just say 'command args' in a private message.
                // If the channel is the sender's nick, the message is being sent in a private message.
                if !command_matched && message.mask.nick() == Some(channel) {
                    let command = message.args[1].slice_from(1); // slice_from(1) to remove the `:` at the beginning of privmsg content.
                    let args = message.args.slice_from(2).iter().map(|s| s.clone()).collect::<Vec<String>>();
                    self.dispatch_command(&plugins, command, channel, args, &message.mask);
                }
            }
        }
    }

    fn dispatch_command(&self, plugins: &sync::RWLockReadGuard<client::PluginRegister>, command: &str, channel: &str, args: Vec<String>, mask: &irc::IrcMask) {
        if let Some(list) = plugins.commands.get(&command.to_ascii_lowercase()) {
            let command_event = events::CommandTransport::new(channel, args, mask);
            for closure in list.iter() {
                self.execute(PluginTask::Command((closure.clone(), command_event.clone())));
            }
        }
    }

    fn execute(&self, task: PluginTask) {
        self.workers_out.send(task);
    }
}

/// TODO: Better name for this
enum PluginTask {
    Command((sync::Arc<client::CommandListener>, events::CommandTransport)),
    Message((sync::Arc<client::MessageListener>, events::MessageTransport)),
    Ctcp((sync::Arc<client::CtcpListener>, events::CtcpTransport)),
}

impl PluginTask {
    fn execute(self, interface: &interface::IrcInterface) {
        match self {
            PluginTask::Command((closure, event)) => {
                closure.call((&events::CommandEvent::new(interface, &event),));
            },
            PluginTask::Message((closure, event)) => {
                closure.call((&events::MessageEvent::new(interface, &event),));
            },
            PluginTask::Ctcp((closure, event)) => {
                closure.call((&events::CtcpEvent::new(interface, &event),));
            },
        }
    }
}

struct PluginExecutor {
    interface: interface::IrcInterface,
    data_in: sync::Arc<sync::Mutex<Receiver<PluginTask>>>,
    logger: fern::ArcLogger,
    active: bool,
}

impl PluginExecutor {
    fn new(interface: interface::IrcInterface, data_in: sync::Arc<sync::Mutex<Receiver<PluginTask>>>, logger: fern::ArcLogger) -> PluginExecutor {
        return PluginExecutor {
            interface: interface,
            data_in: data_in,
            logger: logger,
            active: true,
        };
    }

    fn worker_loop(&mut self) {
        loop {
            let message = {
                // Only lock jobs for the time it takes
                // to get a job, not run it.
                let lock = self.data_in.lock().unwrap();
                lock.recv_opt()
            };
            match message {
                Ok(next) => next.execute(&self.interface),
                Err(()) => {
                    self.active = false;
                    break;
                }
            };
        }
    }

    fn start_worker_thread(mut self) {
        thread::Builder::new().name("worker_thread".to_string()).spawn(move || {
            fern::local::set_thread_logger(self.logger.clone());
            self.worker_loop();
        }).detach();
    }
}

impl Drop for PluginExecutor {
    fn drop(&mut self) {
        if self.active {
            warning!("Worker panicked!");
            PluginExecutor::new(self.interface.clone(), self.data_in.clone(), self.logger.clone()).start_worker_thread();
        }
    }
}
<|endoftext|>


<commit_before>extern crate num_traits;

use std::f64;
use std::cmp::Ordering;
use std::str::FromStr;
use std::ops::{Add, Mul};
use self::num_traits::{One, Zero};

#[derive(PartialOrd, Debug, Default, Clone, Copy)]
pub struct LogProb {
    value: f64
}

impl LogProb {
    pub fn new(value: f64) -> Result<Self, String> {
        if 0.0 <= value && value <= 1.0 {
            Ok(LogProb { value: - value.log2() })
        } else {
            Err(format!("Value {} is not a probability (i.e. not in the interval [0,1]).", value))
        }
    }

    pub fn probability(&self) -> f64 {
        (- &self.value).exp2()
    }
}

/// Flipped instance to reflect that a probability of ``0.25`` is "more expensive" than ``0.5``.
impl Ord for LogProb {
    fn cmp(&self, other: &Self) -> Ordering {
        match self.partial_cmp(&other) {
            Some(ordering) => ordering,
            None => {
                if self.value.is_nan() {
                    if other.value.is_nan() {
                        Ordering::Equal
                    } else {
                        Ordering::Greater
                    }
                } else {
                    Ordering::Less
                }
            }
        }
    }
}

impl PartialEq for LogProb {
    fn eq(&self, other: &Self) -> bool {
        if self.value.is_nan() {
            if other.value.is_nan() { true } else { false }
        } else if other.value.is_nan() {
            false
        } else {
            self.value == other.value
        }
    }
}

impl Eq for LogProb {}

impl Mul for LogProb {
    type Output = Self;

    fn mul(self, other: Self) -> Self {
        LogProb { value: self.value.add(other.value) }
    }
}

impl Add for LogProb {
    type Output = Self;

    fn add(self, other: Self) -> Self {
        match LogProb::new(self.probability().add(other.probability())) {
            Ok(p) => p,
            Err(e) => panic!(e)
        }
    }
}

impl Zero for LogProb {
    fn zero() -> LogProb {
        LogProb { value: f64::INFINITY }
    }

    fn is_zero(&self) -> bool {
        self.value.is_infinite()
    }
}

impl One for LogProb {
    fn one() -> LogProb {
        LogProb { value: 0.0 }
    }
}

impl FromStr for LogProb {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.parse() {
            Ok(p) => LogProb::new(p),
            Err(e) => Err(e.to_string())
        }
    }
}

<commit_msg>deleted unnecessary copy of log_prob<commit_after><|endoftext|>