<commit_before>use std::cmp::Ordering;
use std::fmt;

use crate::conversion::ToPyObject;
use crate::err::{self, PyErr, PyResult};
use crate::ffi;
use crate::objects::{PyObject, PyInt, PyLong, PyFloat};
use crate::python::{Python, PythonObject, ToPythonPointer};

use super::ObjectProtocol;

/// Operations on numeric objects
pub trait NumberProtocol: ObjectProtocol {
    /// Perform addition (self + other)
    ///
    /// Invokes the `__add__` magic-method
    #[inline]
    fn add(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Add(self.as_ptr(), other))
        })
    }
    /// Perform subtraction (self - other)
    ///
    /// Invokes the `__sub__` magic-method
    #[inline]
    fn subtract(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Subtract(self.as_ptr(), other))
        })
    }
    /// Perform multiplication (self * other)
    ///
    /// Invokes the `__mul__` magic-method
    #[inline]
    fn multiply(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Multiply(self.as_ptr(), other))
        })
    }
    /// Perform matrix multiplication, equivalent to the Python expression `self @ other`
    ///
    /// Invokes the `__matmul__` magic-method
    ///
    /// This was added in Python 3.5, and will unconditionally
    /// throw an exception on any version before that.
    ///
    /// See [PEP 0456](https://www.python.org/dev/peps/pep-0465/) for details.
    #[inline]
    fn matrix_multiply(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        #[cfg(not(any(feature = "python3-4", feature = "python2-sys")))] {
            other.with_borrowed_ptr(py, |other| unsafe {
                err::result_from_owned_ptr(py, ffi::PyNumber_MatrixMultiply(self.as_ptr(), other))
            })
        }
        #[cfg(any(feature = "python3-4", feature = "python2-sys"))] {

            drop(other);
            Err(crate::PyErr::new::<crate::exc::TypeError, _>(
                py,
                "Matrix multiplication is unsupported before Python 3.5"
            ))
        }
    }
    /// Perform exponentiation, equivalent to the Python expression `self ** other`,
    /// or the two-argument form of the builtin method pow: `pow(self, other)`
    ///
    /// Invokes the `__pow__` magic-method
    ///
    /// See also [NumberProtocol::pow_mod].
    #[inline]
    fn pow(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        self.pow_mod(py, other, py.None())
    }
    /// Perform exponentiation modulo an integer,
    /// mathematically equivalent to `self ** other % mod`
    /// but computed much more efficiently.
    ///
    /// Equivalent to invoking the three-argument form
    /// of the builtin `pow` method: `pow(self, other, z)`
    ///
    /// Invoking with a `None` for modulo is equivalent to
    /// the regular power operation.
    ///
    /// Invokes the `__pow__` magic-method
    #[inline]
    fn pow_mod(&self, py: Python, exp: impl ToPyObject, z: impl ToPyObject) -> PyResult<PyObject> {
        exp.with_borrowed_ptr(py, |exp| {
            z.with_borrowed_ptr(py, |z| unsafe {
                err::result_from_owned_ptr(py, ffi::PyNumber_Power(self.as_ptr(), exp, z))
            })
        })
    }
    /// Perform the "true division" operation,
    /// equivalent to the Python expression `self / other`,
    ///
    /// Invokes the `__truediv__` magic-method.
    #[inline]
    fn true_div(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_TrueDivide(self.as_ptr(), other))
        })
    }
    /// Perform the "floor division" operation,
    /// equivalent to the Python expression `self // other`,
    ///
    /// This method was added in Python 3.
    /// If compiling against Python 2, it unconditional throws an error.
    ///
    /// Invokes the `__floordiv__` magic-method.
    #[inline]
    fn div_floor(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_FloorDivide(self.as_ptr(), other))
        })
    }
    /// Return the remainder of dividing `self` by `other`,
    /// equivalent to the Python expression `self % other`
    ///
    /// Invokes the `__mod__` magic-method.
    #[inline]
    fn remainder(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Remainder(self.as_ptr(), other))
        })
    }
    /// Perform combined division and modulo,
    /// equivalent to the builtin method `divmod(self, other)`
    ///
    /// Invokes the `__divmod__` magic-method.
    #[inline]
    fn div_mod(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Divmod(self.as_ptr(), other))
        })
    }
    /// Perform the negation of self (-self)
    ///
    /// Invokes the `__neg__` magic-method.
    #[inline]
    fn negative(&self, py: Python) -> PyResult<PyObject> {
        unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Negative(self.as_ptr()))
        }
    }
    /// Invoke the 'positive' operation, equivalent to the
    /// Python expression `+self`
    ///
    /// Invokes the `__pos__` magic-method
    #[inline]
    fn positive(&self, py: Python) -> PyResult<PyObject> {
        unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Positive(self.as_ptr()))
        }
    }
    /// Return the absolute value of self,
    /// equivalent to calling the builtin function `abs`
    ///
    /// Invokes the `__abs__` magic-method.
    #[inline]
    fn absolute(&self, py: Python) -> PyResult<PyObject> {
        unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Absolute(self.as_ptr()))
        }
    }
    /// Perform the bitwise negation of self,
    /// equivalent to the Python expression `~self`
    ///
    /// Invokes the `__invert__` magic-method
    #[inline]
    fn bitwise_invert(&self, py: Python) -> PyResult<PyObject> {
        unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Invert(self.as_ptr()))
        }
    }
    /// Shift this value to the left by the specified number of bits,
    /// equivalent to the Python expression `self << bits`
    ///
    /// Invokes the `__lshift__` magic-method
    #[inline]
    fn left_shift(&self, py: Python, bits: impl ToPyObject) -> PyResult<PyObject> {
        bits.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Lshift(self.as_ptr(), other))
        })
    }
    /// Shift this value to the right by the specified number of bits,
    /// equivalent to the Python expression `self >> bits`
    ///
    /// Invokes the `__rshift__` magic-method
    #[inline]
    fn right_shift(&self, py: Python, bits: impl ToPyObject) -> PyResult<PyObject> {
        bits.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Rshift(self.as_ptr(), other))
        })
    }
    /// Perform the "bitwise and" of `self & other`
    ///
    /// Invokes the `__and__` magic-method.
    #[inline]
    fn bitwise_and(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_And(self.as_ptr(), other))
        })
    }
    /// Perform the "bitwise exclusive or",
    /// equivalent to Python expression `self ^ other`
    ///
    /// Invokes the `__xor__` magic-method.
    #[inline]
    fn bitwise_xor(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Xor(self.as_ptr(), other))
        })
    }
    /// Perform the "bitwise or" of `self | other`
    ///
    /// Invokes the `__or__` magic-method.
    #[inline]
    fn bitwise_or(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Or(self.as_ptr(), other))
        })
    }
    /// Convert this object to an integer,
    /// equivalent to the builtin function `int(self)`
    ///
    /// Invokes the `__int__` magic-method.
    ///
    /// Throws an exception if unable to perform
    /// the conversion.
    #[inline]
    fn as_int(&self, py: Python) -> PyResult<PyLong> {
        let obj = unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Long(self.as_ptr()))?
        };
        Ok(obj.cast_into::<PyLong>(py)?)
    }
    /// Convert this object to a float,
    /// equivalent to the builtin function `float(self)`
    ///
    /// Invokes the `__float__` magic-method.
    ///
    /// Throws an exception if unable to perform
    /// the conversion.
    #[inline]
    fn as_float(&self, py: Python) -> PyResult<PyFloat> {
        let obj = unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Float(self.as_ptr()))?
        };
        Ok(obj.cast_into::<PyFloat>(py)?)
    }
    /// Losslessly convert this object to an integer index,
    /// as if calling `operator.index()`
    ///
    /// The presence of this method indicates
    /// this object is an integer type.
    ///
    /// Calls the `__index__` magic-method.
    ///
    /// See also: [Documentation on the corresponding magic-method](https://docs.python.org/3/reference/datamodel.html?highlight=__index__#object.__index__)
    #[inline]
    fn as_int_index(&self, py: Python) -> PyResult<PyLong> {
        let obj = unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Index(self.as_ptr()))?
        };
        Ok(obj.cast_into::<PyLong>(py)?)
    }
}

impl NumberProtocol for PyObject {}

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

    #[test]
    fn addition() {
        let guard = Python::acquire_gil();
        let py = guard.python();
        let i1 = (5i32).to_py_object(py).into_object();
        let i2 = (12i32).to_py_object(py).into_object();
        let actual_res = i1.add(py, i2).unwrap();
        let expected_res = (17i32).to_py_object(py).into_object();
        assert_eq!(
            actual_res.compare(py, expected_res).unwrap(),
            Ordering::Equal
        );
    }

    py_class!(class DummyMatMul |py| {
        data number: i32;
        def __new__(_cls, arg: i32) -> PyResult<DummyMatMul> {
            DummyMatMul::create_instance(py, arg)
        }
        def __matmul__(left, other) -> PyResult<PyObject> {
            // Do a dummy operation that can be easily tested
            left.cast_as::<Self>(py)?.number(py)
                .to_py_object(py)
                .into_object()
                .multiply(py, other)?
                .add(py, 3)
        }
    });

    #[test]
    #[cfg_attr(any(feature = "python3-4", feature = "python2-sys"), should_panic)]
    fn matrix_multiply() {
        let guard = Python::acquire_gil();
        let py = guard.python();
        let first = DummyMatMul::create_instance(py, 5).unwrap().into_object();
        let seven = (7i32).to_py_object(py).into_object();
        let actual_res = first.matrix_multiply(py, seven).unwrap();
        // 5 * 7 + 3 => 38
        let expected_res = (38i32).to_py_object(py).into_object();
        assert_eq!(
            actual_res.compare(py, expected_res).unwrap(),
            Ordering::Equal
        );
    }
}
<commit_msg>Name NumberProtocol naming more consistent<commit_after>use std::cmp::Ordering;
use std::fmt;

use crate::conversion::ToPyObject;
use crate::err::{self, PyErr, PyResult};
use crate::ffi;
use crate::objects::{PyObject, PyInt, PyLong, PyFloat};
use crate::python::{Python, PythonObject, ToPythonPointer};

use super::ObjectProtocol;

/// Operations on numeric objects
pub trait NumberProtocol: ObjectProtocol {
    /// Perform addition (self + other)
    ///
    /// Invokes the `__add__` magic-method
    #[inline]
    fn add(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Add(self.as_ptr(), other))
        })
    }
    /// Perform subtraction (self - other)
    ///
    /// Invokes the `__sub__` magic-method
    #[inline]
    fn subtract(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Subtract(self.as_ptr(), other))
        })
    }
    /// Perform multiplication (self * other)
    ///
    /// Invokes the `__mul__` magic-method
    #[inline]
    fn multiply(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Multiply(self.as_ptr(), other))
        })
    }
    /// Perform matrix multiplication, equivalent to the Python expression `self @ other`
    ///
    /// Invokes the `__matmul__` magic-method
    ///
    /// This was added in Python 3.5, and will unconditionally
    /// throw an exception on any version before that.
    ///
    /// See [PEP 0456](https://www.python.org/dev/peps/pep-0465/) for details.
    #[inline]
    fn matrix_multiply(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        #[cfg(not(any(feature = "python3-4", feature = "python2-sys")))] {
            other.with_borrowed_ptr(py, |other| unsafe {
                err::result_from_owned_ptr(py, ffi::PyNumber_MatrixMultiply(self.as_ptr(), other))
            })
        }
        #[cfg(any(feature = "python3-4", feature = "python2-sys"))] {

            drop(other);
            Err(crate::PyErr::new::<crate::exc::TypeError, _>(
                py,
                "Matrix multiplication is unsupported before Python 3.5"
            ))
        }
    }
    /// Perform exponentiation, equivalent to the Python expression `self ** other`,
    /// or the two-argument form of the builtin method pow: `pow(self, other)`
    ///
    /// Invokes the `__pow__` magic-method
    ///
    /// See also [NumberProtocol::power_modulo].
    #[inline]
    fn power(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        self.power_modulo(py, other, py.None())
    }
    /// Perform exponentiation modulo an integer,
    /// mathematically equivalent to `self ** other % mod`
    /// but computed much more efficiently.
    ///
    /// Equivalent to invoking the three-argument form
    /// of the builtin `pow` method: `pow(self, other, z)`
    ///
    /// Invoking with a `None` for modulo is equivalent to
    /// the regular power operation.
    ///
    /// Invokes the `__pow__` magic-method
    #[inline]
    fn power_modulo(&self, py: Python, exp: impl ToPyObject, z: impl ToPyObject) -> PyResult<PyObject> {
        exp.with_borrowed_ptr(py, |exp| {
            z.with_borrowed_ptr(py, |z| unsafe {
                err::result_from_owned_ptr(py, ffi::PyNumber_Power(self.as_ptr(), exp, z))
            })
        })
    }
    /// Perform the "true division" operation,
    /// equivalent to the Python expression `self / other`,
    ///
    /// Invokes the `__truediv__` magic-method.
    #[inline]
    fn true_divide(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_TrueDivide(self.as_ptr(), other))
        })
    }
    /// Perform the "floor division" operation,
    /// equivalent to the Python expression `self // other`,
    ///
    /// This method was added in Python 3.
    /// If compiling against Python 2, it unconditional throws an error.
    ///
    /// Invokes the `__floordiv__` magic-method.
    #[inline]
    fn floor_divide(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_FloorDivide(self.as_ptr(), other))
        })
    }
    /// Return the remainder of dividing `self` by `other`,
    /// equivalent to the Python expression `self % other`
    ///
    /// Invokes the `__mod__` magic-method.
    #[inline]
    fn modulo(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Remainder(self.as_ptr(), other))
        })
    }
    /// Perform combined division and modulo,
    /// equivalent to the builtin method `divmod(self, other)`
    ///
    /// Invokes the `__divmod__` magic-method.
    #[inline]
    fn div_mod(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Divmod(self.as_ptr(), other))
        })
    }
    /// Perform the negation of self (-self)
    ///
    /// Invokes the `__neg__` magic-method.
    #[inline]
    fn negative(&self, py: Python) -> PyResult<PyObject> {
        unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Negative(self.as_ptr()))
        }
    }
    /// Invoke the 'positive' operation, equivalent to the
    /// Python expression `+self`
    ///
    /// Invokes the `__pos__` magic-method
    #[inline]
    fn positive(&self, py: Python) -> PyResult<PyObject> {
        unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Positive(self.as_ptr()))
        }
    }
    /// Return the absolute value of self,
    /// equivalent to calling the builtin function `abs`
    ///
    /// Invokes the `__abs__` magic-method.
    #[inline]
    fn absolute(&self, py: Python) -> PyResult<PyObject> {
        unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Absolute(self.as_ptr()))
        }
    }
    /// Perform the bitwise negation of self,
    /// equivalent to the Python expression `~self`
    ///
    /// Invokes the `__invert__` magic-method
    #[inline]
    fn bitwise_invert(&self, py: Python) -> PyResult<PyObject> {
        unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Invert(self.as_ptr()))
        }
    }
    /// Shift this value to the left by the specified number of bits,
    /// equivalent to the Python expression `self << bits`
    ///
    /// Invokes the `__lshift__` magic-method
    #[inline]
    fn left_shift(&self, py: Python, bits: impl ToPyObject) -> PyResult<PyObject> {
        bits.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Lshift(self.as_ptr(), other))
        })
    }
    /// Shift this value to the right by the specified number of bits,
    /// equivalent to the Python expression `self >> bits`
    ///
    /// Invokes the `__rshift__` magic-method
    #[inline]
    fn right_shift(&self, py: Python, bits: impl ToPyObject) -> PyResult<PyObject> {
        bits.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Rshift(self.as_ptr(), other))
        })
    }
    /// Perform the "bitwise and" of `self & other`
    ///
    /// Invokes the `__and__` magic-method.
    #[inline]
    fn bitwise_and(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_And(self.as_ptr(), other))
        })
    }
    /// Perform the "bitwise exclusive or",
    /// equivalent to Python expression `self ^ other`
    ///
    /// Invokes the `__xor__` magic-method.
    #[inline]
    fn bitwise_xor(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Xor(self.as_ptr(), other))
        })
    }
    /// Perform the "bitwise or" of `self | other`
    ///
    /// Invokes the `__or__` magic-method.
    #[inline]
    fn bitwise_or(&self, py: Python, other: impl ToPyObject) -> PyResult<PyObject> {
        other.with_borrowed_ptr(py, |other| unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Or(self.as_ptr(), other))
        })
    }
    /// Convert this object to an integer,
    /// equivalent to the builtin function `int(self)`
    ///
    /// Invokes the `__int__` magic-method.
    ///
    /// Throws an exception if unable to perform
    /// the conversion.
    #[inline]
    fn as_int(&self, py: Python) -> PyResult<PyLong> {
        let obj = unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Long(self.as_ptr()))?
        };
        Ok(obj.cast_into::<PyLong>(py)?)
    }
    /// Convert this object to a float,
    /// equivalent to the builtin function `float(self)`
    ///
    /// Invokes the `__float__` magic-method.
    ///
    /// Throws an exception if unable to perform
    /// the conversion.
    #[inline]
    fn as_float(&self, py: Python) -> PyResult<PyFloat> {
        let obj = unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Float(self.as_ptr()))?
        };
        Ok(obj.cast_into::<PyFloat>(py)?)
    }
    /// Losslessly convert this object to an integer index,
    /// as if calling `operator.index()`
    ///
    /// The presence of this method indicates
    /// this object is an integer type.
    ///
    /// Calls the `__index__` magic-method.
    ///
    /// See also: [Documentation on the corresponding magic-method](https://docs.python.org/3/reference/datamodel.html?highlight=__index__#object.__index__)
    #[inline]
    fn as_int_index(&self, py: Python) -> PyResult<PyLong> {
        let obj = unsafe {
            err::result_from_owned_ptr(py, ffi::PyNumber_Index(self.as_ptr()))?
        };
        Ok(obj.cast_into::<PyLong>(py)?)
    }
}

impl NumberProtocol for PyObject {}

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

    #[test]
    fn addition() {
        let guard = Python::acquire_gil();
        let py = guard.python();
        let i1 = (5i32).to_py_object(py).into_object();
        let i2 = (12i32).to_py_object(py).into_object();
        let actual_res = i1.add(py, i2).unwrap();
        let expected_res = (17i32).to_py_object(py).into_object();
        assert_eq!(
            actual_res.compare(py, expected_res).unwrap(),
            Ordering::Equal
        );
    }

    py_class!(class DummyMatMul |py| {
        data number: i32;
        def __new__(_cls, arg: i32) -> PyResult<DummyMatMul> {
            DummyMatMul::create_instance(py, arg)
        }
        def __matmul__(left, other) -> PyResult<PyObject> {
            // Do a dummy operation that can be easily tested
            left.cast_as::<Self>(py)?.number(py)
                .to_py_object(py)
                .into_object()
                .multiply(py, other)?
                .add(py, 3)
        }
    });

    #[test]
    #[cfg_attr(any(feature = "python3-4", feature = "python2-sys"), should_panic)]
    fn matrix_multiply() {
        let guard = Python::acquire_gil();
        let py = guard.python();
        let first = DummyMatMul::create_instance(py, 5).unwrap().into_object();
        let seven = (7i32).to_py_object(py).into_object();
        let actual_res = first.matrix_multiply(py, seven).unwrap();
        // 5 * 7 + 3 => 38
        let expected_res = (38i32).to_py_object(py).into_object();
        assert_eq!(
            actual_res.compare(py, expected_res).unwrap(),
            Ordering::Equal
        );
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 Intecture Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

//! Apt package provider

use {Command, CommandResult, Host};
use Result;
use super::*;

pub struct Apt;

impl Provider for Apt {
    fn get_providers(&self) -> Providers {
        Providers::Apt
    }

    fn is_active(&self, host: &mut Host) -> Result<bool> {
        let cmd = Command::new("which apt-get");
        let result = try!(cmd.exec(host));

        Ok(result.exit_code == 0)
    }

    fn is_installed(&self, host: &mut Host, name: &str) -> Result<bool> {
        let cmd = Command::new(&format!("dpkg -l {}", name));
        let result = try!(cmd.exec(host));

        Ok(result.exit_code == 0)
    }

    fn install(&self, host: &mut Host, name: &str) -> Result<CommandResult> {
        let cmd = Command::new(&format!("apt-get -y install {}", name));
        cmd.exec(host)
    }

    fn uninstall(&self, host: &mut Host, name: &str) -> Result<CommandResult> {
        let cmd = Command::new(&format!("apt-get -y remove {}", name));
        cmd.exec(host)
    }
}
<commit_msg>Fixed bug detecting installed Apt packages<commit_after>// Copyright 2015 Intecture Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

//! Apt package provider

use {Command, CommandResult, Host};
use Result;
use super::*;

pub struct Apt;

impl Provider for Apt {
    fn get_providers(&self) -> Providers {
        Providers::Apt
    }

    fn is_active(&self, host: &mut Host) -> Result<bool> {
        let cmd = Command::new("which apt-get");
        let result = try!(cmd.exec(host));

        Ok(result.exit_code == 0)
    }

    fn is_installed(&self, host: &mut Host, name: &str) -> Result<bool> {
        let cmd = Command::new(&format!("dpkg --get-selections | grep -E \"{}\\s+install$\"", name));
        let result = try!(cmd.exec(host));

        Ok(result.exit_code == 0)
    }

    fn install(&self, host: &mut Host, name: &str) -> Result<CommandResult> {
        let cmd = Command::new(&format!("apt-get -y install {}", name));
        cmd.exec(host)
    }

    fn uninstall(&self, host: &mut Host, name: &str) -> Result<CommandResult> {
        let cmd = Command::new(&format!("apt-get -y remove {}", name));
        cmd.exec(host)
    }
}
<|endoftext|>


<commit_before>use utils::lru_cache::LRUCache;
use dovahkiin::types::custom_types::id::Id;
use dovahkiin::types::custom_types::map::Map;
use std::cell::RefCell;
use ram::cell::Cell;
use std::cell::Ref;
use std::cell::RefMut;
use index::btree::*;
use bifrost::utils::async_locks::{RwLock, RwLockWriteGuard, RwLockReadGuard, Mutex};
use std::ops::Deref;
use std::ops::DerefMut;
use core::borrow::BorrowMut;
use dovahkiin::types::value::ToValue;
use itertools::Itertools;
use std::collections::BTreeMap;
use std::mem;

pub type ExtNodeCacheMap = Mutex<LRUCache<Id, RwLock<ExtNode>>>;
pub type ExtNodeCachedMut = RwLockWriteGuard<ExtNode>;
pub type ExtNodeCachedImmute = RwLockReadGuard<ExtNode>;

lazy_static! {
    static ref KEYS_KEY_HASH : u64 = key_hash("keys");
    static ref NEXT_PAGE_KEY_HASH : u64 = key_hash("next");
    static ref PREV_PAGE_KEY_HASH : u64 = key_hash("prev");
    static ref PAGE_SCHEMA_ID: u32 = key_hash("BTREE_SCHEMA_ID") as u32;
}

#[derive(Clone)]
pub struct ExtNode {
    pub id: Id,
    pub keys: EntryKeySlice,
    pub next: Id,
    pub prev: Id,
    pub len: usize,
    pub version: u64,
    pub removed: bool,
}

pub struct ExtNodeSplit {
    pub node_2: ExtNode,
    pub keys_1_len: usize
}

impl ExtNode {
    pub fn new(id: &Id) -> ExtNode {
        ExtNode {
            id: *id,
            keys: EntryKeySlice::init(),
            next: Id::unit_id(),
            prev: Id::unit_id(),
            len: 0,
            version: 0,
            removed: false
        }
    }
    pub fn from_cell(cell: Cell) -> Self {
        let keys = &cell.data[*KEYS_KEY_HASH];
        let keys_len = keys.len().unwrap();
        let mut key_slice = EntryKeySlice::init();
        let mut key_count = 0;
        let next = cell.data[*NEXT_PAGE_KEY_HASH].Id().unwrap();
        let prev = cell.data[*PREV_PAGE_KEY_HASH].Id().unwrap();
        for (i, key_val) in keys.iter_value().unwrap().enumerate() {
            let key = if let Value::PrimArray(PrimitiveArray::U8(ref array)) = key_val {
                EntryKey::from_slice(array.as_slice())
            } else { panic!("invalid entry") };
            key_slice[i] = key;
            key_count += 1;
        }
        ExtNode {
            id: cell.id(),
            keys: key_slice,
            next: *next,
            prev: *prev,
            len: key_count,
            version: cell.header.version,
            removed: false
        }
    }
    pub fn to_cell(&self) -> Cell {
        let mut value = Value::Map(Map::new());
        value[*NEXT_PAGE_KEY_HASH] = Value::Id(self.next);
        value[*KEYS_KEY_HASH] = self
            .keys[..self.len]
            .iter()
            .map(|key| {
                key.as_slice().to_vec().value()
            })
            .collect_vec()
            .value();
        Cell::new_with_id(*PAGE_SCHEMA_ID, &self.id, value)
    }
    pub fn remove(&mut self, pos: usize) {
        let cached_len = self.len;
        self.keys.remove_at(pos, cached_len);
        self.len -= 1;
    }
    pub fn insert(&mut self, key: EntryKey, pos: usize, tree: &BPlusTree) -> Option<(Node, Option<EntryKey>)> {
        let mut cached = self;
        let cached_len = cached.len;
        if cached_len + 1 >= NUM_KEYS {
            // need to split
            let pivot = cached_len / 2;
            let split = {
                let cached_next = *&cached.next;
                let cached_id = *&cached.id;
                let new_page_id = tree.new_page_id();
                cached.next = new_page_id;
                let mut keys_1 = &mut cached.keys;
                let mut keys_2 = keys_1.split_at_pivot(pivot, cached_len);
                let mut keys_1_len = pivot;
                let mut keys_2_len = cached_len - pivot;
                insert_into_split(
                    key,
                    keys_1, &mut keys_2,
                    &mut keys_1_len, &mut keys_2_len,
                    pos, pivot);
                ExtNodeSplit {
                    keys_1_len,
                    node_2: ExtNode {
                        id: new_page_id,
                        keys: keys_2,
                        next: cached_next,
                        prev: cached_id,
                        len: keys_2_len,
                        version: 0,
                        removed: false
                    }
                }
            };
            cached.next = split.node_2.id;
            cached.len = split.keys_1_len;
            return Some((Node::External(cached.id), None));

        } else {
            cached.keys.insert_at(key, pos, cached_len);
            return None;
        }
    }
    pub fn merge_with(&mut self, right: &mut Self) {
        let self_len = self.len;
        let new_len = self.len + right.len;
        assert!(new_len <= self.keys.len());
        for i in self.len .. new_len {
            self.keys[i] = mem::replace(&mut right.keys[i - self_len], Default::default());
        }
    }
}

pub fn rearrange_empty_extnode(node: &ExtNode, bz: &mut CacheBufferZone) -> Id {
    let prev = node.prev;
    let next = node.next;
    if !prev.is_unit_id() {
        let mut prev_node = bz.get_for_mut(&prev);
        prev_node.next = next;
    }
    if !next.is_unit_id() {
        let mut next_node = bz.get_for_mut(&next);
        next_node.prev = prev;
    }
    return node.id;
}

#[derive(Clone)]
enum CacheGuardHolder {
    Read(ExtNodeCachedImmute),
    Write(ExtNodeCachedMut)
}

pub struct CacheBufferZone<'a> {
    tree: &'a BPlusTree,
    guards: BTreeMap<Id, CacheGuardHolder>,
    data: BTreeMap<Id, Option<RefCell<ExtNode>>>
}

impl <'a> CacheBufferZone <'a> {
    pub fn new(tree: &BPlusTree) -> CacheBufferZone {
        CacheBufferZone {
            tree,
            guards: BTreeMap::new(),
            data: BTreeMap::new()
        }
    }

    pub fn get(&mut self, id: &Id) -> Ref<ExtNode> {
        {
            match self.data.get(id) {
                Some(Some(ref data)) => return data.borrow(),
                Some(None) => panic!(),
                _ => {}
            }
        }
        let guard = self.tree.get_ext_node_cached(id);
        self.data.insert(*id, Some(RefCell::new((&*guard).clone())));
        let holder = CacheGuardHolder::Read(guard);
        self.guards.insert(*id, holder);
        return self.get(id);
    }

    pub fn get_for_mut(&mut self, id: &Id) -> RefMut<ExtNode> {
        {
            match self.data.get(id) {
                Some(Some(ref data)) => {
                    if let CacheGuardHolder::Read(_) = self.guards.get(id).unwrap() {
                        panic!("Mutating readonly buffered cache");
                    }
                    return data.borrow_mut();
                },
                Some(None) => panic!(),
                _ => {}
            }
        }
        let guard = self.tree.get_mut_ext_node_cached(id);
        self.data.insert(*id, Some(RefCell::new((&*guard).clone())));
        let holder = CacheGuardHolder::Write(guard);
        self.guards.insert(*id, holder);
        self.get_for_mut(id)
    }

    pub fn update(&mut self, id: &Id, data: ExtNode) {
        self.data.insert(*id, Some(RefCell::new(data)));
    }
    pub fn delete(&mut self, id: &Id) {
        self.data.insert(*id, None);
    }

    pub fn flush(mut self) {
        for (id, data) in self.data {
            if let Some(node) = data {
                let mut holder = self.guards.get_mut(&id).unwrap();
                if let &mut CacheGuardHolder::Write(ref mut guard) = holder {
                    **guard = node.into_inner()
                } else { panic!() }
            } else {
                unimplemented!();
            }
        }
    }
}<commit_msg>owning reference<commit_after>use utils::lru_cache::LRUCache;
use dovahkiin::types::custom_types::id::Id;
use dovahkiin::types::custom_types::map::Map;
use std::cell::RefCell;
use ram::cell::Cell;
use std::cell::Ref;
use std::cell::RefMut;
use index::btree::*;
use bifrost::utils::async_locks::{RwLock, RwLockWriteGuard, RwLockReadGuard, Mutex};
use std::ops::Deref;
use std::ops::DerefMut;
use core::borrow::BorrowMut;
use dovahkiin::types::value::ToValue;
use itertools::Itertools;
use std::collections::BTreeMap;
use std::mem;
use std::rc::Rc;
use owning_ref::{OwningRefMut, OwningRef, RcRef};

pub type ExtNodeCacheMap = Mutex<LRUCache<Id, RwLock<ExtNode>>>;
pub type ExtNodeCachedMut = RwLockWriteGuard<ExtNode>;
pub type ExtNodeCachedImmute = RwLockReadGuard<ExtNode>;

lazy_static! {
    static ref KEYS_KEY_HASH : u64 = key_hash("keys");
    static ref NEXT_PAGE_KEY_HASH : u64 = key_hash("next");
    static ref PREV_PAGE_KEY_HASH : u64 = key_hash("prev");
    static ref PAGE_SCHEMA_ID: u32 = key_hash("BTREE_SCHEMA_ID") as u32;
}

#[derive(Clone)]
pub struct ExtNode {
    pub id: Id,
    pub keys: EntryKeySlice,
    pub next: Id,
    pub prev: Id,
    pub len: usize,
    pub version: u64,
    pub removed: bool,
}

pub struct ExtNodeSplit {
    pub node_2: ExtNode,
    pub keys_1_len: usize
}

impl ExtNode {
    pub fn new(id: &Id) -> ExtNode {
        ExtNode {
            id: *id,
            keys: EntryKeySlice::init(),
            next: Id::unit_id(),
            prev: Id::unit_id(),
            len: 0,
            version: 0,
            removed: false
        }
    }
    pub fn from_cell(cell: Cell) -> Self {
        let keys = &cell.data[*KEYS_KEY_HASH];
        let keys_len = keys.len().unwrap();
        let mut key_slice = EntryKeySlice::init();
        let mut key_count = 0;
        let next = cell.data[*NEXT_PAGE_KEY_HASH].Id().unwrap();
        let prev = cell.data[*PREV_PAGE_KEY_HASH].Id().unwrap();
        for (i, key_val) in keys.iter_value().unwrap().enumerate() {
            let key = if let Value::PrimArray(PrimitiveArray::U8(ref array)) = key_val {
                EntryKey::from_slice(array.as_slice())
            } else { panic!("invalid entry") };
            key_slice[i] = key;
            key_count += 1;
        }
        ExtNode {
            id: cell.id(),
            keys: key_slice,
            next: *next,
            prev: *prev,
            len: key_count,
            version: cell.header.version,
            removed: false
        }
    }
    pub fn to_cell(&self) -> Cell {
        let mut value = Value::Map(Map::new());
        value[*NEXT_PAGE_KEY_HASH] = Value::Id(self.next);
        value[*KEYS_KEY_HASH] = self
            .keys[..self.len]
            .iter()
            .map(|key| {
                key.as_slice().to_vec().value()
            })
            .collect_vec()
            .value();
        Cell::new_with_id(*PAGE_SCHEMA_ID, &self.id, value)
    }
    pub fn remove(&mut self, pos: usize) {
        let cached_len = self.len;
        self.keys.remove_at(pos, cached_len);
        self.len -= 1;
    }
    pub fn insert(&mut self, key: EntryKey, pos: usize, tree: &BPlusTree) -> Option<(Node, Option<EntryKey>)> {
        let mut cached = self;
        let cached_len = cached.len;
        if cached_len + 1 >= NUM_KEYS {
            // need to split
            let pivot = cached_len / 2;
            let split = {
                let cached_next = *&cached.next;
                let cached_id = *&cached.id;
                let new_page_id = tree.new_page_id();
                cached.next = new_page_id;
                let mut keys_1 = &mut cached.keys;
                let mut keys_2 = keys_1.split_at_pivot(pivot, cached_len);
                let mut keys_1_len = pivot;
                let mut keys_2_len = cached_len - pivot;
                insert_into_split(
                    key,
                    keys_1, &mut keys_2,
                    &mut keys_1_len, &mut keys_2_len,
                    pos, pivot);
                ExtNodeSplit {
                    keys_1_len,
                    node_2: ExtNode {
                        id: new_page_id,
                        keys: keys_2,
                        next: cached_next,
                        prev: cached_id,
                        len: keys_2_len,
                        version: 0,
                        removed: false
                    }
                }
            };
            cached.next = split.node_2.id;
            cached.len = split.keys_1_len;
            return Some((Node::External(cached.id), None));

        } else {
            cached.keys.insert_at(key, pos, cached_len);
            return None;
        }
    }
    pub fn merge_with(&mut self, right: &mut Self) {
        let self_len = self.len;
        let new_len = self.len + right.len;
        assert!(new_len <= self.keys.len());
        for i in self.len .. new_len {
            self.keys[i] = mem::replace(&mut right.keys[i - self_len], Default::default());
        }
    }
}

pub fn rearrange_empty_extnode(node: &ExtNode, bz: &mut CacheBufferZone) -> Id {
    let prev = node.prev;
    let next = node.next;
    if !prev.is_unit_id() {
        let mut prev_node = bz.get_for_mut(&prev);
        prev_node.next = next;
    }
    if !next.is_unit_id() {
        let mut next_node = bz.get_for_mut(&next);
        next_node.prev = prev;
    }
    return node.id;
}

#[derive(Clone)]
enum CacheGuardHolder {
    Read(ExtNodeCachedImmute),
    Write(ExtNodeCachedMut)
}

type NodeRcRefCell = Rc<RefCell<ExtNode>>;

pub struct RcNodeRef {
    owner: NodeRcRefCell,
    reference: *const ExtNode
}

pub struct RcNodeRefMut {
    owner: NodeRcRefCell,
    reference: *mut ExtNode
}

impl RcNodeRef {
    pub fn new(rc: &NodeRcRefCell) -> RcNodeRef {
        let rc = rc.clone();
        RcNodeRef {
            reference: &*(*rc).borrow(),
            owner: rc
        }
    }
}

impl RcNodeRefMut {
    pub fn new(rc: &NodeRcRefCell) -> RcNodeRefMut {
        let rc = rc.clone();
        RcNodeRefMut {
            reference: &mut *(*rc).borrow_mut(),
            owner: rc
        }
    }
}

impl Deref for RcNodeRef {
    type Target = ExtNode;

    fn deref(&self) -> &'_ <Self as Deref>::Target {
        unsafe {
            &* self.reference
        }
    }
}

impl Deref for RcNodeRefMut {
    type Target = ExtNode;

    fn deref(&self) -> &'_ <Self as Deref>::Target {
        unsafe {
            &* self.reference
        }
    }
}

impl DerefMut for RcNodeRefMut {
    fn deref_mut(&mut self) -> &'_ mut <Self as Deref>::Target {
        unsafe {
            &mut* self.reference
        }
    }
}

pub struct CacheBufferZone<'a> {
    tree: &'a BPlusTree,
    guards: RefCell<BTreeMap<Id, CacheGuardHolder>>,
    data: RefCell<BTreeMap<Id, Option<NodeRcRefCell>>>
}

impl <'a> CacheBufferZone <'a> {
    pub fn new(tree: &BPlusTree) -> CacheBufferZone {
        CacheBufferZone {
            tree,
            guards: RefCell::new(BTreeMap::new()),
            data: RefCell::new(BTreeMap::new())
        }
    }

    pub fn get(&self, id: &Id) -> RcNodeRef {
        {
            let data = self.data.borrow();
            match data.get(id) {
                Some(Some(ref data)) =>
                    return RcNodeRef::new(data),
                Some(None) => panic!(),
                _ => {}
            }
        }
        {
            let guard = self.tree.get_ext_node_cached(id);
            let mut data = self.data.borrow_mut();
            data.insert(*id, Some(Rc::new(RefCell::new((&*guard).clone()))));
            let holder = CacheGuardHolder::Read(guard);
            let mut guards = self.guards.borrow_mut();
            guards.insert(*id, holder);
        }
        return self.get(id);
    }

    pub fn get_for_mut(&self, id: &Id) -> RcNodeRefMut {
        {
            let data = self.data.borrow();
            match data.get(id) {
                Some(Some(ref data)) => {
                    let guards = self.guards.borrow();
                    if let CacheGuardHolder::Read(_) = guards.get(id).unwrap() {
                        panic!("Mutating readonly buffered cache");
                    }
                    return RcNodeRefMut::new(data);
                },
                Some(None) => panic!(),
                _ => {}
            }
        }
        {
            let guard = self.tree.get_mut_ext_node_cached(id);
            let mut data = self.data.borrow_mut();
            data.borrow_mut().insert(*id, Some(Rc::new(RefCell::new((&*guard).clone()))));
            let holder = CacheGuardHolder::Write(guard);
            let mut guards = self.guards.borrow_mut();
            guards.insert(*id, holder);
        }
        self.get_for_mut(id)
    }

    pub fn update(&mut self, id: &Id, node: ExtNode) {
        let mut data = self.data.borrow_mut();
        data.insert(*id, Some(Rc::new(RefCell::new(node))));
    }
    pub fn delete(&mut self, id: &Id) {
        let mut data = self.data.borrow_mut();
        data.insert(*id, None);
    }

    pub fn flush(self) {
        let mut guards = self.guards.into_inner();
        for (id, data) in self.data.into_inner() {
            if let Some(node) = data {
                let mut holder = guards.get_mut(&id).unwrap();
                if let &mut CacheGuardHolder::Write(ref mut guard) = holder {
                    **guard = node.into_inner()
                } else { panic!() }
            } else {
                unimplemented!();
            }
        }
    }
}<|endoftext|>


<commit_before>use schema::Schema;
use schema::Document;
use indexer::SegmentSerializer;
use core::Index;
use core::SerializableSegment;
use core::Segment;
use std::thread::JoinHandle;
use indexer::SegmentWriter;
use std::clone::Clone;
use std::io;
use std::path::Path;
use std::thread;
use std::collections::HashSet;
use indexer::merger::IndexMerger;
use core::SegmentId;
use datastruct::stacker::Heap;
use std::mem::swap;
use chan;
use directory::WritePtr;

use Result;
use Error;

// Size of the margin for the heap. A segment is closed when the remaining memory
// in the heap goes below MARGIN_IN_BYTES.
pub const MARGIN_IN_BYTES: u32 = 10_000_000u32;

// We impose the memory per thread to be at least 30 MB.
pub const HEAP_SIZE_LIMIT: u32 = MARGIN_IN_BYTES * 3u32;

// Add document will block if the number of docs waiting in the queue to be indexed reaches PIPELINE_MAX_SIZE_IN_DOCS
const PIPELINE_MAX_SIZE_IN_DOCS: usize = 10_000;

pub const LOCKFILE_NAME: &'static str = ".tantivy-indexer.lock";

type DocumentSender = chan::Sender<Document>;
type DocumentReceiver = chan::Receiver<Document>;

type NewSegmentSender = chan::Sender<Result<(SegmentId, usize)>>;
type NewSegmentReceiver = chan::Receiver<Result<(SegmentId, usize)>>;

/// `IndexWriter` is the user entry-point to add document to an index.
///
/// It manages a small number of indexing thread, as well as a shared
/// indexing queue.
/// Each indexing thread builds its own independant `Segment`, via
/// a `SegmentWriter` object.
pub struct IndexWriter {
	index: Index,
	heap_size_in_bytes_per_thread: usize,
	workers_join_handle: Vec<JoinHandle<()>>,
	segment_ready_sender: NewSegmentSender,
	segment_ready_receiver: NewSegmentReceiver,
	document_receiver: DocumentReceiver,
	document_sender: DocumentSender,
	num_threads: usize,
	  docstamp: u64,
    lockfile: WritePtr
}

fn index_documents(heap: &mut Heap,
				   segment: Segment,
				   schema: &Schema,
				   document_iterator: &mut Iterator<Item=Document>) -> Result<usize> {
	heap.clear();
	let mut segment_writer = try!(SegmentWriter::for_segment(heap, segment, &schema));
	for doc in document_iterator {
		try!(segment_writer.add_document(&doc, &schema));
		if segment_writer.is_buffer_full() {
			info!("Buffer limit reached, flushing segment with maxdoc={}.", segment_writer.max_doc());
			break;
		}
	}
	let num_docs = segment_writer.max_doc() as usize;
	try!(segment_writer.finalize());
	Ok(num_docs)
}

impl IndexWriter {

	/// Spawns a new worker thread for indexing.
	/// The thread consumes documents from the pipeline.
	///
	fn add_indexing_worker(&mut self,) -> Result<()> {
		let index = self.index.clone();
		let schema = self.index.schema();
		let segment_ready_sender_clone = self.segment_ready_sender.clone();
		let document_receiver_clone = self.document_receiver.clone();

		let mut heap = Heap::with_capacity(self.heap_size_in_bytes_per_thread);
		let join_handle: JoinHandle<()> = thread::spawn(move || {
			loop {
				let segment = index.new_segment();
				let segment_id = segment.id();
				let mut document_iterator = document_receiver_clone
					.clone()
					.into_iter()
					.peekable();
				// the peeking here is to avoid
				// creating a new segment's files
				// if no document are available.
				if document_iterator.peek().is_some() {
					let index_result = index_documents(&mut heap, segment, &schema, &mut document_iterator)
						.map(|num_docs| (segment_id, num_docs));
					segment_ready_sender_clone.send(index_result);
				}
				else {
					return;
				}
			}
		});
		self.workers_join_handle.push(join_handle);
		Ok(())
	}

	/// Open a new index writer
	///
	/// num_threads tells the number of indexing worker that
	/// should work at the same time.
	pub fn open(index: &Index,
				num_threads: usize,
				heap_size_in_bytes_per_thread: usize) -> Result<IndexWriter> {
		if heap_size_in_bytes_per_thread <= HEAP_SIZE_LIMIT as usize {
			panic!(format!("The heap size per thread needs to be at least {}.", HEAP_SIZE_LIMIT));
		}

      let mut cloned_index = index.clone();
      let lockfile_path = Path::new(LOCKFILE_NAME);

      let lf = try!(cloned_index.directory_mut().open_write(lockfile_path));
		let (document_sender, document_receiver): (DocumentSender, DocumentReceiver) = chan::sync(PIPELINE_MAX_SIZE_IN_DOCS);
		let (segment_ready_sender, segment_ready_receiver): (NewSegmentSender, NewSegmentReceiver) = chan::async();
		let mut index_writer = IndexWriter {
			heap_size_in_bytes_per_thread: heap_size_in_bytes_per_thread,
			index: cloned_index,
			segment_ready_receiver: segment_ready_receiver,
			segment_ready_sender: segment_ready_sender,
			document_receiver: document_receiver,
			document_sender: document_sender,
			workers_join_handle: Vec::new(),
			num_threads: num_threads,
			docstamp: try!(index.docstamp()),
        lockfile: lf
		};
		try!(index_writer.start_workers());
		Ok(index_writer)
	}

	fn start_workers(&mut self,) -> Result<()> {
		for _ in 0 .. self.num_threads {
			try!(self.add_indexing_worker());
		}
		Ok(())
	}

	/// Merges a given list of segments
	pub fn merge(&mut self, segments: &[Segment]) -> Result<()> {
		let schema = self.index.schema();
		let merger = try!(IndexMerger::open(schema, segments));
		let mut merged_segment = self.index.new_segment();
		let segment_serializer = try!(SegmentSerializer::for_segment(&mut merged_segment));
		try!(merger.write(segment_serializer));
		let merged_segment_ids: HashSet<SegmentId> = segments.iter().map(|segment| segment.id()).collect();
		try!(self.index.publish_merge_segment(merged_segment_ids, merged_segment.id()));
		Ok(())
	}

	/// Closes the current document channel send.
	/// and replace all the channels by new ones.
	///
	/// The current workers will keep on indexing
	/// the pending document and stop
	/// when no documents are remaining.
	///
	/// Returns the former segment_ready channel.
	fn recreate_channels(&mut self,) -> (DocumentReceiver, chan::Receiver<Result<(SegmentId, usize)>>) {
		let (mut document_sender, mut document_receiver): (DocumentSender, DocumentReceiver) = chan::sync(PIPELINE_MAX_SIZE_IN_DOCS);
		let (mut segment_ready_sender, mut segment_ready_receiver): (NewSegmentSender, NewSegmentReceiver) = chan::async();
		swap(&mut self.document_sender, &mut document_sender);
		swap(&mut self.document_receiver, &mut document_receiver);
		swap(&mut self.segment_ready_sender, &mut segment_ready_sender);
		swap(&mut self.segment_ready_receiver, &mut segment_ready_receiver);
		(document_receiver, segment_ready_receiver)
	}

	/// Rollback to the last commit
	///
	/// This cancels all of the update that
	/// happened before after the last commit.
	/// After calling rollback, the index is in the same
	/// state as it was after the last commit.
	///
	/// The docstamp at the last commit is returned.
	pub fn rollback(&mut self,) -> Result<u64> {

		// we cannot drop segment ready receiver yet
		// as it would block the workers.
		let (document_receiver, mut _segment_ready_receiver) = self.recreate_channels();

		// consumes the document receiver pipeline
		// worker don't need to index the pending documents.
		for _ in document_receiver {};

		let mut former_workers_join_handle = Vec::new();
		swap(&mut former_workers_join_handle, &mut self.workers_join_handle);

		// wait for all the worker to finish their work
		// (it should be fast since we consumed all pending documents)
		for worker_handle in former_workers_join_handle {
			try!(worker_handle
				.join()
				.map_err(|e| Error::ErrorInThread(format!("{:?}", e)))
			);
			// add a new worker for the next generation.
			try!(self.add_indexing_worker());
		}

		// reset the docstamp to what it was before
		self.docstamp = try!(self.index.docstamp());
		Ok(self.docstamp)
	}

	/// Commits all of the pending changes
	///
	/// A call to commit blocks.
	/// After it returns, all of the document that
	/// were added since the last commit are published
	/// and persisted.
	///
	/// In case of a crash or an hardware failure (as
	/// long as the hard disk is spared), it will be possible
	/// to resume indexing from this point.
	///
	/// Commit returns the `docstamp` of the last document
	/// that made it in the commit.
	///
	pub fn commit(&mut self,) -> Result<u64> {

		let (document_receiver, segment_ready_receiver) = self.recreate_channels();
		drop(document_receiver);

		// Docstamp of the last document in this commit.
		let commit_docstamp = self.docstamp;

		let mut former_workers_join_handle = Vec::new();
		swap(&mut former_workers_join_handle, &mut self.workers_join_handle);

		for worker_handle in former_workers_join_handle {
			try!(worker_handle
				.join()
				.map_err(|e| Error::ErrorInThread(format!("{:?}", e)))
			);
			// add a new worker for the next generation.
			try!(self.add_indexing_worker());
		}

		let segment_ids_and_size: Vec<(SegmentId, usize)> = try!(
			segment_ready_receiver
				.into_iter()
				.collect()
		);

		let segment_ids: Vec<SegmentId> = segment_ids_and_size
			.iter()
			.map(|&(segment_id, _num_docs)| segment_id)
			.collect();

		try!(self.index.publish_segments(&segment_ids, commit_docstamp));

		Ok(commit_docstamp)
	}

	/// Adds a document.
	///
	/// If the indexing pipeline is full, this call may block.
	///
	/// The docstamp is an increasing `u64` that can
	/// be used by the client to align commits with its own
	/// document queue.
	///
	/// Currently it represents the number of documents that
	/// have been added since the creation of the index.
	pub fn add_document(&mut self, doc: Document) -> io::Result<u64> {
		self.document_sender.send(doc);
		self.docstamp += 1;
		Ok(self.docstamp)
	}

}

#[cfg(test)]
mod tests {

	use schema::{self, Document};
	use Index;
	use Term;

	#[test]
	fn test_commit_and_rollback() {
		let mut schema_builder = schema::SchemaBuilder::default();
		let text_field = schema_builder.add_text_field("text", schema::TEXT);
		let index = Index::create_in_ram(schema_builder.build());

		let num_docs_containing = |s: &str| {
			let searcher = index.searcher();
			let term_a = Term::from_field_text(text_field, s);
			searcher.doc_freq(&term_a)
		};

		{
			// writing the segment
			let mut index_writer = index.writer_with_num_threads(3, 40_000_000).unwrap();
			{
				let mut doc = Document::default();
				doc.add_text(text_field, "a");
				index_writer.add_document(doc).unwrap();
			}
			assert_eq!(index_writer.rollback().unwrap(), 0u64);
			assert_eq!(num_docs_containing("a"), 0);

			{
				let mut doc = Document::default();
				doc.add_text(text_field, "b");
				index_writer.add_document(doc).unwrap();
			}
			{
				let mut doc = Document::default();
				doc.add_text(text_field, "c");
				index_writer.add_document(doc).unwrap();
			}
			assert_eq!(index_writer.commit().unwrap(), 2u64);

			assert_eq!(num_docs_containing("a"), 0);
			assert_eq!(num_docs_containing("b"), 1);
			assert_eq!(num_docs_containing("c"), 1);
		}
		index.searcher();
	}

}
<commit_msg>delete lockfile on drop IndexWriter<commit_after>use schema::Schema;
use schema::Document;
use indexer::SegmentSerializer;
use core::Index;
use core::SerializableSegment;
use core::Segment;
use std::thread::JoinHandle;
use indexer::SegmentWriter;
use std::clone::Clone;
use std::io;
use std::path::Path;
use std::thread;
use std::collections::HashSet;
use indexer::merger::IndexMerger;
use core::SegmentId;
use datastruct::stacker::Heap;
use std::mem::swap;
use chan;
use directory::WritePtr;

use Result;
use Error;

// Size of the margin for the heap. A segment is closed when the remaining memory
// in the heap goes below MARGIN_IN_BYTES.
pub const MARGIN_IN_BYTES: u32 = 10_000_000u32;

// We impose the memory per thread to be at least 30 MB.
pub const HEAP_SIZE_LIMIT: u32 = MARGIN_IN_BYTES * 3u32;

// Add document will block if the number of docs waiting in the queue to be indexed reaches PIPELINE_MAX_SIZE_IN_DOCS
const PIPELINE_MAX_SIZE_IN_DOCS: usize = 10_000;

pub const LOCKFILE_NAME: &'static str = ".tantivy-indexer.lock";

type DocumentSender = chan::Sender<Document>;
type DocumentReceiver = chan::Receiver<Document>;

type NewSegmentSender = chan::Sender<Result<(SegmentId, usize)>>;
type NewSegmentReceiver = chan::Receiver<Result<(SegmentId, usize)>>;

/// `IndexWriter` is the user entry-point to add document to an index.
///
/// It manages a small number of indexing thread, as well as a shared
/// indexing queue.
/// Each indexing thread builds its own independant `Segment`, via
/// a `SegmentWriter` object.
pub struct IndexWriter {
	index: Index,
	heap_size_in_bytes_per_thread: usize,
	workers_join_handle: Vec<JoinHandle<()>>,
	segment_ready_sender: NewSegmentSender,
	segment_ready_receiver: NewSegmentReceiver,
	document_receiver: DocumentReceiver,
	document_sender: DocumentSender,
	num_threads: usize,
	  docstamp: u64,
    lockfile: WritePtr
}

fn index_documents(heap: &mut Heap,
				   segment: Segment,
				   schema: &Schema,
				   document_iterator: &mut Iterator<Item=Document>) -> Result<usize> {
	heap.clear();
	let mut segment_writer = try!(SegmentWriter::for_segment(heap, segment, &schema));
	for doc in document_iterator {
		try!(segment_writer.add_document(&doc, &schema));
		if segment_writer.is_buffer_full() {
			info!("Buffer limit reached, flushing segment with maxdoc={}.", segment_writer.max_doc());
			break;
		}
	}
	let num_docs = segment_writer.max_doc() as usize;
	try!(segment_writer.finalize());
	Ok(num_docs)
}

impl Drop for IndexWriter {
    fn drop(&mut self) {
        let lockfile_path = Path::new(LOCKFILE_NAME);
        match self.index.directory_mut().delete(lockfile_path) {
            Ok(_) => (),
            Err(_) => ()
        }
    }
}
impl IndexWriter {

	/// Spawns a new worker thread for indexing.
	/// The thread consumes documents from the pipeline.
	///
	fn add_indexing_worker(&mut self,) -> Result<()> {
		let index = self.index.clone();
		let schema = self.index.schema();
		let segment_ready_sender_clone = self.segment_ready_sender.clone();
		let document_receiver_clone = self.document_receiver.clone();

		let mut heap = Heap::with_capacity(self.heap_size_in_bytes_per_thread);
		let join_handle: JoinHandle<()> = thread::spawn(move || {
			loop {
				let segment = index.new_segment();
				let segment_id = segment.id();
				let mut document_iterator = document_receiver_clone
					.clone()
					.into_iter()
					.peekable();
				// the peeking here is to avoid
				// creating a new segment's files
				// if no document are available.
				if document_iterator.peek().is_some() {
					let index_result = index_documents(&mut heap, segment, &schema, &mut document_iterator)
						.map(|num_docs| (segment_id, num_docs));
					segment_ready_sender_clone.send(index_result);
				}
				else {
					return;
				}
			}
		});
		self.workers_join_handle.push(join_handle);
		Ok(())
	}

	/// Open a new index writer
	///
	/// num_threads tells the number of indexing worker that
	/// should work at the same time.
	pub fn open(index: &Index,
				num_threads: usize,
				heap_size_in_bytes_per_thread: usize) -> Result<IndexWriter> {
		if heap_size_in_bytes_per_thread <= HEAP_SIZE_LIMIT as usize {
			panic!(format!("The heap size per thread needs to be at least {}.", HEAP_SIZE_LIMIT));
		}

      let mut cloned_index = index.clone();
      let lockfile_path = Path::new(LOCKFILE_NAME);

      let lf = try!(cloned_index.directory_mut().open_write(lockfile_path));
		let (document_sender, document_receiver): (DocumentSender, DocumentReceiver) = chan::sync(PIPELINE_MAX_SIZE_IN_DOCS);
		let (segment_ready_sender, segment_ready_receiver): (NewSegmentSender, NewSegmentReceiver) = chan::async();
		let mut index_writer = IndexWriter {
			heap_size_in_bytes_per_thread: heap_size_in_bytes_per_thread,
			index: cloned_index,
			segment_ready_receiver: segment_ready_receiver,
			segment_ready_sender: segment_ready_sender,
			document_receiver: document_receiver,
			document_sender: document_sender,
			workers_join_handle: Vec::new(),
			num_threads: num_threads,
			docstamp: try!(index.docstamp()),
        lockfile: lf
		};
		try!(index_writer.start_workers());
		Ok(index_writer)
	}

	fn start_workers(&mut self,) -> Result<()> {
		for _ in 0 .. self.num_threads {
			try!(self.add_indexing_worker());
		}
		Ok(())
	}

	/// Merges a given list of segments
	pub fn merge(&mut self, segments: &[Segment]) -> Result<()> {
		let schema = self.index.schema();
		let merger = try!(IndexMerger::open(schema, segments));
		let mut merged_segment = self.index.new_segment();
		let segment_serializer = try!(SegmentSerializer::for_segment(&mut merged_segment));
		try!(merger.write(segment_serializer));
		let merged_segment_ids: HashSet<SegmentId> = segments.iter().map(|segment| segment.id()).collect();
		try!(self.index.publish_merge_segment(merged_segment_ids, merged_segment.id()));
		Ok(())
	}

	/// Closes the current document channel send.
	/// and replace all the channels by new ones.
	///
	/// The current workers will keep on indexing
	/// the pending document and stop
	/// when no documents are remaining.
	///
	/// Returns the former segment_ready channel.
	fn recreate_channels(&mut self,) -> (DocumentReceiver, chan::Receiver<Result<(SegmentId, usize)>>) {
		let (mut document_sender, mut document_receiver): (DocumentSender, DocumentReceiver) = chan::sync(PIPELINE_MAX_SIZE_IN_DOCS);
		let (mut segment_ready_sender, mut segment_ready_receiver): (NewSegmentSender, NewSegmentReceiver) = chan::async();
		swap(&mut self.document_sender, &mut document_sender);
		swap(&mut self.document_receiver, &mut document_receiver);
		swap(&mut self.segment_ready_sender, &mut segment_ready_sender);
		swap(&mut self.segment_ready_receiver, &mut segment_ready_receiver);
		(document_receiver, segment_ready_receiver)
	}

	/// Rollback to the last commit
	///
	/// This cancels all of the update that
	/// happened before after the last commit.
	/// After calling rollback, the index is in the same
	/// state as it was after the last commit.
	///
	/// The docstamp at the last commit is returned.
	pub fn rollback(&mut self,) -> Result<u64> {

		// we cannot drop segment ready receiver yet
		// as it would block the workers.
		let (document_receiver, mut _segment_ready_receiver) = self.recreate_channels();

		// consumes the document receiver pipeline
		// worker don't need to index the pending documents.
		for _ in document_receiver {};

		let mut former_workers_join_handle = Vec::new();
		swap(&mut former_workers_join_handle, &mut self.workers_join_handle);

		// wait for all the worker to finish their work
		// (it should be fast since we consumed all pending documents)
		for worker_handle in former_workers_join_handle {
			try!(worker_handle
				.join()
				.map_err(|e| Error::ErrorInThread(format!("{:?}", e)))
			);
			// add a new worker for the next generation.
			try!(self.add_indexing_worker());
		}

		// reset the docstamp to what it was before
		self.docstamp = try!(self.index.docstamp());
		Ok(self.docstamp)
	}

	/// Commits all of the pending changes
	///
	/// A call to commit blocks.
	/// After it returns, all of the document that
	/// were added since the last commit are published
	/// and persisted.
	///
	/// In case of a crash or an hardware failure (as
	/// long as the hard disk is spared), it will be possible
	/// to resume indexing from this point.
	///
	/// Commit returns the `docstamp` of the last document
	/// that made it in the commit.
	///
	pub fn commit(&mut self,) -> Result<u64> {

		let (document_receiver, segment_ready_receiver) = self.recreate_channels();
		drop(document_receiver);

		// Docstamp of the last document in this commit.
		let commit_docstamp = self.docstamp;

		let mut former_workers_join_handle = Vec::new();
		swap(&mut former_workers_join_handle, &mut self.workers_join_handle);

		for worker_handle in former_workers_join_handle {
			try!(worker_handle
				.join()
				.map_err(|e| Error::ErrorInThread(format!("{:?}", e)))
			);
			// add a new worker for the next generation.
			try!(self.add_indexing_worker());
		}

		let segment_ids_and_size: Vec<(SegmentId, usize)> = try!(
			segment_ready_receiver
				.into_iter()
				.collect()
		);

		let segment_ids: Vec<SegmentId> = segment_ids_and_size
			.iter()
			.map(|&(segment_id, _num_docs)| segment_id)
			.collect();

		try!(self.index.publish_segments(&segment_ids, commit_docstamp));

		Ok(commit_docstamp)
	}

	/// Adds a document.
	///
	/// If the indexing pipeline is full, this call may block.
	///
	/// The docstamp is an increasing `u64` that can
	/// be used by the client to align commits with its own
	/// document queue.
	///
	/// Currently it represents the number of documents that
	/// have been added since the creation of the index.
	pub fn add_document(&mut self, doc: Document) -> io::Result<u64> {
		self.document_sender.send(doc);
		self.docstamp += 1;
		Ok(self.docstamp)
	}

}

#[cfg(test)]
mod tests {

	use schema::{self, Document};
	use Index;
	use Term;

	#[test]
	fn test_commit_and_rollback() {
		let mut schema_builder = schema::SchemaBuilder::default();
		let text_field = schema_builder.add_text_field("text", schema::TEXT);
		let index = Index::create_in_ram(schema_builder.build());

		let num_docs_containing = |s: &str| {
			let searcher = index.searcher();
			let term_a = Term::from_field_text(text_field, s);
			searcher.doc_freq(&term_a)
		};

		{
			// writing the segment
			let mut index_writer = index.writer_with_num_threads(3, 40_000_000).unwrap();
			{
				let mut doc = Document::default();
				doc.add_text(text_field, "a");
				index_writer.add_document(doc).unwrap();
			}
			assert_eq!(index_writer.rollback().unwrap(), 0u64);
			assert_eq!(num_docs_containing("a"), 0);

			{
				let mut doc = Document::default();
				doc.add_text(text_field, "b");
				index_writer.add_document(doc).unwrap();
			}
			{
				let mut doc = Document::default();
				doc.add_text(text_field, "c");
				index_writer.add_document(doc).unwrap();
			}
			assert_eq!(index_writer.commit().unwrap(), 2u64);

			assert_eq!(num_docs_containing("a"), 0);
			assert_eq!(num_docs_containing("b"), 1);
			assert_eq!(num_docs_containing("c"), 1);
		}
		index.searcher();
	}

}
<|endoftext|>


<commit_before>use super::expr::Expression;
use super::{CompileError, is_truthy};
use super::scope::Scope;
use super::sum::Sum;

#[cfg(test)]
use super::lexer;
#[cfg(test)]
use super::{TRUE, FALSE};

// Something that can be called with arguments given in the scope
pub trait Function {
	fn call(&self, scope: &Scope) -> Result<f32, CompileError>;
}

macro_rules! bind_function(
	( $name:ident, $func:ident ( $($arg:ident),* ) ) => (
		struct $name;
		impl $name {
			fn new() -> $name {
				$name
			}
		}
		impl Function for $name {
			#[allow(unused_variables)] // Compiler complains that scope is not used on functions with no args.
			fn call(&self, scope: &Scope) -> Result<f32, CompileError> {
				Ok($func($(
					// Insert each argument from the scope
					match scope.var_id(stringify!($arg)) {
						Some(id) => scope.get_var(id).unwrap(),

						None => return Err(CompileError {
							msg: format!("function `{}` requires argument `{}` but it is not defined in scope", stringify!($func), stringify!($arg)),
							pos: None }),
					}
				),*))
			}
		}
	);
)

fn sin(a: f32) -> f32 {
	a.sin()
}

bind_function!(SinFunction, sin(a))

// Always returns a specific constant
struct ConstFunction {
	val: f32,
}
impl ConstFunction {
	fn new(val: f32) -> ConstFunction {
		ConstFunction {
			val: val,
		}
	}
}
impl Function for ConstFunction {
	fn call(&self, _: &Scope) -> Result<f32, CompileError> {
		Ok(self.val)
	}
}

// Wraps an expression from the expr module.
struct ExprFunction<'a> {
	expr: &'a Expression<'a>,
}
impl<'a> ExprFunction<'a> {
	fn new(expr: &'a Expression<'a>) -> ExprFunction<'a> {
		ExprFunction {
			expr: expr,
		}
	}
}
impl<'a> Function for ExprFunction<'a> {
	fn call(&self, scope: &Scope) -> Result<f32, CompileError> {
		self.expr.eval(scope)
	}
}

// Returns the result of calling the function if the condition is truthy, else 0
struct CondFunction<'a> {
	cond: Box<Function + 'a>,
	func: Box<Function + 'a>,
}
impl<'a> CondFunction<'a> {
	fn new(cond: Box<Function + 'a>, func: Box<Function + 'a>) -> CondFunction<'a> {
		CondFunction {
			cond: cond,
			func: func,
		}
	}
}
impl<'a> Function for CondFunction<'a> {
	fn call(&self, scope: &Scope) -> Result<f32, CompileError> {
		if is_truthy(try!(self.cond.call(scope))) {
			self.func.call(scope)
		} else {
			Ok(0_f32)
		}
	}
}

// A SumFunction represents the sum of a bunch of other functions. Functions defined in syntizer
// are internally represented as this.
struct SumFunction<'a> {
	sum: &'a Sum<'a>
}
impl<'a> SumFunction<'a> {
	pub fn new(sum: &'a Sum<'a>) -> SumFunction<'a> {
		SumFunction {
			sum: sum,
		}
	}
}
impl<'a> Function for SumFunction<'a> {
	fn call(&self, scope: &Scope) -> Result<f32, CompileError> {
		self.sum.eval(scope)
	}
}

#[test]
fn const_function_test() {
	let f = ConstFunction::new(42_f32);
	assert!(f.call(&Scope::new()).unwrap() == 42_f32);
}

#[test]
fn expr_function_test() {
	let mut s = Scope::new();
	s.define_var("a", 9_f32);
	let e = Expression::new(lexer::tokenize("5*3+a").unwrap().as_slice(), &s).unwrap();
	let f = ExprFunction::new(&e);
	assert!(f.call(&s).unwrap() == 24_f32);
}

#[test]
fn sin_test() {
	let a = 2_f32;
	let mut s = Scope::new();
	s.define_var("a", a);
	let f = SinFunction::new();
	assert!(f.call(&s).unwrap() == a.sin());
}

#[test]
fn cond_test() {
	let cond_truthy = ConstFunction::new(TRUE);
	let cond_falsey = ConstFunction::new(FALSE);
	let expr = ConstFunction::new(42_f32);
	let f_truthy = CondFunction::new(box cond_truthy, box expr);
	let f_falsey = CondFunction::new(box cond_falsey, box expr);
	let s = Scope::new();
	assert!(f_truthy.call(&s).unwrap() == 42_f32);
	assert!(f_falsey.call(&s).unwrap() == 0_f32);
}

#[test]
fn sum_test() {
	let sum = Sum::new(vec![
		box ConstFunction::new(1_f32) as Box<Function>,
		box ConstFunction::new(2_f32) as Box<Function>,
		box ConstFunction::new(3_f32) as Box<Function>,
		box ConstFunction::new(4_f32) as Box<Function>]);
	let f = SumFunction::new(&sum);
	let s = Scope::new();
	assert!(f.call(&s).unwrap() == 10_f32);
}
<commit_msg>Make bind_function! support default args and convert some assert! to assert_eq!<commit_after>use super::expr::Expression;
use super::{CompileError, is_truthy};
use super::scope::Scope;
use super::sum::Sum;

#[cfg(test)]
use super::lexer;
#[cfg(test)]
use super::{TRUE, FALSE};

// Something that can be called with arguments given in the scope
pub trait Function {
	fn call(&self, scope: &Scope) -> Result<f32, CompileError>;
}

macro_rules! bind_function(
	( $name:ident, $func:ident ( $($arg:ident = $val:expr),* ) ) => (
		struct $name;
		impl $name {
			fn new() -> $name {
				$name
			}
		}
		impl Function for $name {
			#[allow(unused_variables)] // Compiler complains that scope is not used on functions with no args.
			fn call(&self, scope: &Scope) -> Result<f32, CompileError> {
				Ok($func($(
					// Insert each argument from the scope
					match scope.var_id(stringify!($arg)) {
						Some(id) => scope.get_var(id).unwrap(),

						None => match $val {
							Some(val) => val,
							None => return Err(CompileError {
								msg: format!("function `{}` requires argument `{}` but it is not defined in scope", stringify!($func), stringify!($arg)),
								pos: None }),
						}
					}
				),*))
			}
		}
	);
)

fn sin(freq: f32, amp: f32, phase: f32, time: f32) -> f32 {
	(freq*time*Float::two_pi() + phase).sin()*amp
}

bind_function!(SinFunction, sin(freq=None, amp=Some(1_f32), phase=Some(0_f32), time=None))

// Always returns a specific constant
struct ConstFunction {
	val: f32,
}
impl ConstFunction {
	fn new(val: f32) -> ConstFunction {
		ConstFunction {
			val: val,
		}
	}
}
impl Function for ConstFunction {
	fn call(&self, _: &Scope) -> Result<f32, CompileError> {
		Ok(self.val)
	}
}

// Wraps an expression from the expr module.
struct ExprFunction<'a> {
	expr: &'a Expression<'a>,
}
impl<'a> ExprFunction<'a> {
	fn new(expr: &'a Expression<'a>) -> ExprFunction<'a> {
		ExprFunction {
			expr: expr,
		}
	}
}
impl<'a> Function for ExprFunction<'a> {
	fn call(&self, scope: &Scope) -> Result<f32, CompileError> {
		self.expr.eval(scope)
	}
}

// Returns the result of calling the function if the condition is truthy, else 0
struct CondFunction<'a> {
	cond: Box<Function + 'a>,
	func: Box<Function + 'a>,
}
impl<'a> CondFunction<'a> {
	fn new(cond: Box<Function + 'a>, func: Box<Function + 'a>) -> CondFunction<'a> {
		CondFunction {
			cond: cond,
			func: func,
		}
	}
}
impl<'a> Function for CondFunction<'a> {
	fn call(&self, scope: &Scope) -> Result<f32, CompileError> {
		if is_truthy(try!(self.cond.call(scope))) {
			self.func.call(scope)
		} else {
			Ok(0_f32)
		}
	}
}

// A SumFunction represents the sum of a bunch of other functions.
struct SumFunction<'a> {
	sum: &'a Sum<'a>
}
impl<'a> SumFunction<'a> {
	pub fn new(sum: &'a Sum<'a>) -> SumFunction<'a> {
		SumFunction {
			sum: sum,
		}
	}
}
impl<'a> Function for SumFunction<'a> {
	fn call(&self, scope: &Scope) -> Result<f32, CompileError> {
		self.sum.eval(scope)
	}
}

#[test]
fn const_function_test() {
	let f = ConstFunction::new(42_f32);
	assert_eq!(f.call(&Scope::new()).unwrap(), 42_f32);
}

#[test]
fn expr_function_test() {
	let mut s = Scope::new();
	s.define_var("a", 9_f32);
	let e = Expression::new(lexer::tokenize("5*3+a").unwrap().as_slice(), &s).unwrap();
	let f = ExprFunction::new(&e);
	assert_eq!(f.call(&s).unwrap(), 24_f32);
}

#[test]
fn sin_test() {
	let mut s = Scope::new();
	s.define_var("freq", 440_f32);
	s.define_var("time", 0_f32);
	s.define_var("amp", 0.5_f32);
	let f = SinFunction::new();
	assert_eq!(f.call(&s).unwrap(), 0_f32);
}

#[test]
fn cond_test() {
	let cond_truthy = ConstFunction::new(TRUE);
	let cond_falsey = ConstFunction::new(FALSE);
	let expr = ConstFunction::new(42_f32);
	let f_truthy = CondFunction::new(box cond_truthy, box expr);
	let f_falsey = CondFunction::new(box cond_falsey, box expr);
	let s = Scope::new();
	assert_eq!(f_truthy.call(&s).unwrap(), 42_f32);
	assert_eq!(f_falsey.call(&s).unwrap(), 0_f32);
}

#[test]
fn sum_test() {
	let sum = Sum::new(vec![
		box ConstFunction::new(1_f32) as Box<Function>,
		box ConstFunction::new(2_f32) as Box<Function>,
		box ConstFunction::new(3_f32) as Box<Function>,
		box ConstFunction::new(4_f32) as Box<Function>]);
	let f = SumFunction::new(&sum);
	let s = Scope::new();
	assert_eq!(f.call(&s).unwrap(), 10_f32);
}
<|endoftext|>


<commit_before>use std::fmt;

#[derive(Debug, Clone)]
struct Tree {
    enter: Node,
    left: Option<Box<Tree>>,
    right: Option<Box<Tree>>,
}

/*impl fmt::Display for Tree {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "(this: {:#?}, \nleft: {:#?}, \nright: {:#?})",
            self.enter, self.left, self.right
        )
    }
}*/

//for red-black tree
#[derive(Debug, Clone, PartialEq)]
enum Color {
    Black,
    Red,
    White,
}

#[derive(Debug, Clone)]
struct Node {
    val: i32,
    color: Color,
}

impl Node {
    fn new(v: &i32) -> Node {
        Node {
            val: *v,
            color: Color::White,
        }
    }

    fn new_black(v: &i32) -> Node {
        let mut n = Node::new(v);
        n.make_black();
        n
    }
    fn new_red(v: &i32) -> Node {
        let mut n = Node::new(v);
        n.make_red();
        n
    }
    fn make_black(&mut self) {
        self.color = Color::Black;
    }

    fn make_red(&mut self) {
        self.color = Color::Red;
    }

    fn change_color(&mut self) {
        if self.color == Color::Black {
            self.color = Color::Red
        } else {
            self.color = Color::Black
        }
    }
}

type ReTr<'a> = Result<&'a mut Tree, String>;

impl Tree {
    fn new(v: &i32) -> Tree {
        let mut root = Node::new(v);
        root.make_black();
        Tree {
            enter: root,
            left: None,
            right: None,
        }
    }

    fn insert(&mut self, v: &i32) {
        let target = if *v > self.enter.val {
            &mut self.right
        } else {
            &mut self.left
        };

        match target {
            Some(next) => next.insert(v),
            None => {
                let mut new_node = Node::new(v);
                new_node.make_red();
                let new_tree = Tree {
                    enter: new_node,
                    left: None,
                    right: None,
                };
                *target = Some(Box::new(new_tree));
            }
        }
    }

    fn look_up_recursive(&mut self, v: &i32) -> ReTr {
        if *v == self.enter.val {
            return Ok(self);
        } else if *v > self.enter.val {
            if let Some(_) = self.right {
                let next = self.right.as_mut().unwrap();
                return next.look_up(v);
            } else {
                Err("not find".to_string())
            }
        } else {
            if let Some(_) = self.left {
                let next = self.left.as_mut().unwrap();
                return next.look_up(v);
            } else {
                Err("not find".to_string())
            }
        }
    }

    fn look_up(&mut self, v: &i32) -> ReTr {
        let mut this: &mut Self = self;
        loop {
            if *v == this.enter.val {
                return Ok(this);
            } else if *v > this.enter.val {
                if let None = this.right {
                    return Err("not find".to_string());
                }
                this = { this }.right.as_mut().unwrap();
            } else {
                if let None = this.left {
                    return Err("not find".to_string());
                }
                this = { this }.left.as_mut().unwrap();
            }
        }
    }

    fn look_up_father(&mut self, v: &i32) -> (ReTr, String) {
        let mut this: &mut Self = self;
        loop {
            if *v == this.enter.val {
                return (Err("no father".to_string()), "this".to_string());
            } else if *v == this.right.as_ref().unwrap().enter.val {
                return (Ok(this), "right".to_string());
            } else if *v == this.left.as_ref().unwrap().enter.val {
                return (Ok(this), "left".to_string());
            } else if *v > this.enter.val {
                if let None = this.right {
                    return (Err("not find".to_string()), "Error".to_string());
                }
                this = { this }.right.as_mut().unwrap();
            } else {
                if let None = this.left {
                    return (Err("not find".to_string()), "Error".to_string());
                }
                this = { this }.left.as_mut().unwrap();
            }
        }
    }

    fn cut_left_tree(&self) -> Option<Box<Self>> {
        self.left.clone()
    }

    fn cut_right_tree(&self) -> Option<Box<Self>> {
        self.right.clone()
    }

    fn left_shift(&mut self, v: &i32) {
        match self.look_up_father(v).0 {
            Ok(r) => {
                let father = r;
                let mut this = father.cut_left_tree().unwrap();
                let mut son = this.cut_right_tree().unwrap();
                let grand_son = son.cut_left_tree().unwrap();

                this.right = Some(grand_son);
                son.left = Some(this);
                father.left = Some(son);
                return;
            }
            _ => (),
        }

        let left = self.cut_left_tree().unwrap();
        let mut new_left = Tree::new(&self.enter.val);
        new_left.left = Some(left);

        let right = self.cut_right_tree().unwrap();
        let right_left = right.cut_left_tree().unwrap();
        let right_right = right.cut_right_tree().unwrap();
        new_left.right = Some(right_left);

        self.enter = Node::new_black(&right.enter.val);

        self.right = Some(right_right);

        self.left = Some(Box::new(new_left));
    }

    fn right_shift(&mut self, v: &i32) {
        match self.look_up_father(v).0 {
            Ok(r) => {
                let father = r;
                let mut this = father.cut_right_tree().unwrap();
                let mut son = this.cut_left_tree().unwrap();
                let grand_son = son.cut_right_tree().unwrap();

                this.left = Some(grand_son);
                son.right = Some(this);
                father.right = Some(son);
                return;
            }
            _ => (),
        }

        let right = self.cut_right_tree().unwrap();
        let left = self.cut_left_tree().unwrap();
        let left_left = left.cut_left_tree().unwrap();
        let left_right = left.cut_left_tree().unwrap();

        let mut new_right = Tree::new(&self.enter.val);
        new_right.left = Some(left_right);
        new_right.right = Some(right);

        self.enter = Node::new_black(&left.enter.val);
        self.left = Some(left_left);
        self.right = Some(Box::new(new_right));
    }
}

fn main() {
    //let mut a = Tree::new(&3);
    //a.insert(&2);
    //a.insert(&1);
    //a.insert(&4);
    //a.insert(&5);

    //println!("{:?}", a);
    //println!("{:?}", a.look_up(&1));
    //println!("{:?}", a.look_up_father(&4));
    //println!("{:?}", a);

    let mut b = Tree::new(&10);
    b.insert(&2);
    b.insert(&1);
    b.insert(&5);
    b.insert(&3);
    b.insert(&6);
    b.insert(&11);

    //println!("{:?}", b);
    //println!("{:?}", b.look_up_father(&2));
    let mut c = b.cut_left_tree().unwrap();
    //println!("{:#?}", c);
    c.left_shift(&2);
    println!("{:#?}", c);

    //let mut c = Tree::new(&2);
    //c.insert(&8);
    //c.insert(&1);
    //c.insert(&5);
    //c.insert(&4);
    //c.insert(&7);
    //c.insert(&9);
    //println!("{:?}", c.look_up_father(&8));
    //c.right_shift(&8);
    //println!("{:#?}", c);
}
<commit_msg>a bit<commit_after>use std::fmt;

#[derive(Debug, Clone)]
struct Tree {
    enter: Node,
    left: Option<Box<Tree>>,
    right: Option<Box<Tree>>,
}

/*impl fmt::Display for Tree {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "(this: {:#?}, \nleft: {:#?}, \nright: {:#?})",
            self.enter, self.left, self.right
        )
    }
}*/

//for red-black tree
#[derive(Debug, Clone, PartialEq)]
enum Color {
    Black,
    Red,
    White,
}

#[derive(Debug, Clone)]
struct Node {
    val: i32,
    color: Color,
}

impl Node {
    fn new(v: &i32) -> Node {
        Node {
            val: *v,
            color: Color::White,
        }
    }

    fn new_black(v: &i32) -> Node {
        let mut n = Node::new(v);
        n.make_black();
        n
    }
    fn new_red(v: &i32) -> Node {
        let mut n = Node::new(v);
        n.make_red();
        n
    }
    fn make_black(&mut self) {
        self.color = Color::Black;
    }

    fn make_red(&mut self) {
        self.color = Color::Red;
    }

    fn change_color(&mut self) {
        if self.color == Color::Black {
            self.color = Color::Red
        } else {
            self.color = Color::Black
        }
    }
}

type ReTr<'a> = Result<&'a mut Tree, String>;

impl Tree {
    fn new(v: &i32) -> Tree {
        let mut root = Node::new(v);
        root.make_black();
        Tree {
            enter: root,
            left: None,
            right: None,
        }
    }

    //fix color at ending of insert method
    fn insert(&mut self, v: &i32) {
        if *v == self.enter.val {
            println!("{}", "duplicate value");
            return;
        }

        {
            //need drop target which hold self ownship..
            //we need this ownship to fix color at endding of this method.
            let target = if *v > self.enter.val {
                &mut self.right
            } else {
                &mut self.left
            };

            match target {
                Some(next) => next.insert(v),
                None => {
                    let mut new_node = Node::new(v);
                    new_node.make_red();
                    let new_tree = Tree {
                        enter: new_node,
                        left: None,
                        right: None,
                    };
                    *target = Some(Box::new(new_tree));
                }
            }
        }

        //self.fix(v);
    }

    fn look_up_recursive(&mut self, v: &i32) -> ReTr {
        if *v == self.enter.val {
            return Ok(self);
        } else if *v > self.enter.val {
            if let Some(_) = self.right {
                let next = self.right.as_mut().unwrap();
                return next.look_up(v);
            } else {
                Err("not find".to_string())
            }
        } else {
            if let Some(_) = self.left {
                let next = self.left.as_mut().unwrap();
                return next.look_up(v);
            } else {
                Err("not find".to_string())
            }
        }
    }

    fn look_up(&mut self, v: &i32) -> ReTr {
        let mut this: &mut Self = self;
        loop {
            if *v == this.enter.val {
                return Ok(this);
            } else if *v > this.enter.val {
                if let None = this.right {
                    return Err("not find".to_string());
                }
                this = { this }.right.as_mut().unwrap();
            } else {
                if let None = this.left {
                    return Err("not find".to_string());
                }
                this = { this }.left.as_mut().unwrap();
            }
        }
    }

    fn look_up_father(&mut self, v: &i32) -> (ReTr, String) {
        let mut this: &mut Self = self;
        loop {
            if *v == this.enter.val {
                return (Err("no father".to_string()), "this".to_string());
            } else if *v == this.right.as_ref().unwrap().enter.val {
                return (Ok(this), "right".to_string());
            } else if *v == this.left.as_ref().unwrap().enter.val {
                return (Ok(this), "left".to_string());
            } else if *v > this.enter.val {
                if let None = this.right {
                    return (Err("not find".to_string()), "Error".to_string());
                }
                this = { this }.right.as_mut().unwrap();
            } else {
                if let None = this.left {
                    return (Err("not find".to_string()), "Error".to_string());
                }
                this = { this }.left.as_mut().unwrap();
            }
        }
    }

    fn cut_left_tree(&self) -> Option<Box<Self>> {
        self.left.clone()
    }

    fn cut_right_tree(&self) -> Option<Box<Self>> {
        self.right.clone()
    }

    fn left_shift(&mut self, v: &i32) {
        match self.look_up_father(v).0 {
            Ok(r) => {
                let father = r;
                let mut this = father.cut_left_tree().unwrap();
                let mut son = this.cut_right_tree().unwrap();
                let grand_son = son.cut_left_tree().unwrap();

                this.right = Some(grand_son);
                son.left = Some(this);
                father.left = Some(son);
                return;
            }
            _ => (),
        }

        let left = self.cut_left_tree().unwrap();
        let mut new_left = Tree::new(&self.enter.val);
        new_left.left = Some(left);

        let right = self.cut_right_tree().unwrap();
        let right_left = right.cut_left_tree().unwrap();
        let right_right = right.cut_right_tree().unwrap();
        new_left.right = Some(right_left);

        self.enter = Node::new_black(&right.enter.val);

        self.right = Some(right_right);

        self.left = Some(Box::new(new_left));
    }

    fn right_shift(&mut self, v: &i32) {
        match self.look_up_father(v).0 {
            Ok(r) => {
                let father = r;
                let mut this = father.cut_right_tree().unwrap();
                let mut son = this.cut_left_tree().unwrap();
                let grand_son = son.cut_right_tree().unwrap();

                this.left = Some(grand_son);
                son.right = Some(this);
                father.right = Some(son);
                return;
            }
            _ => (),
        }

        let right = self.cut_right_tree().unwrap();
        let left = self.cut_left_tree().unwrap();
        let left_left = left.cut_left_tree().unwrap();
        let left_right = left.cut_left_tree().unwrap();

        let mut new_right = Tree::new(&self.enter.val);
        new_right.left = Some(left_right);
        new_right.right = Some(right);

        self.enter = Node::new_black(&left.enter.val);
        self.left = Some(left_left);
        self.right = Some(Box::new(new_right));
    }
}

fn main() {
    //let mut a = Tree::new(&3);
    //a.insert(&2);
    //a.insert(&1);
    //a.insert(&4);
    //a.insert(&5);

    //println!("{:?}", a);
    //println!("{:?}", a.look_up(&1));
    //println!("{:?}", a.look_up_father(&4));
    //println!("{:?}", a);

    let mut b = Tree::new(&10);
    b.insert(&2);
    b.insert(&1);
    b.insert(&5);
    b.insert(&3);
    b.insert(&6);
    b.insert(&11);

    //println!("{:?}", b);
    //println!("{:?}", b.look_up_father(&2));
    let mut c = b.cut_left_tree().unwrap();
    //println!("{:#?}", c);
    c.left_shift(&2);
    println!("{:#?}", c);

    //let mut c = Tree::new(&2);
    //c.insert(&8);
    //c.insert(&1);
    //c.insert(&5);
    //c.insert(&4);
    //c.insert(&7);
    //c.insert(&9);
    //println!("{:?}", c.look_up_father(&8));
    //c.right_shift(&8);
    //println!("{:#?}", c);
}
<|endoftext|>