<commit_before>//! This module contains implementations of the common components used by the
//! various NES sound channels.

#[cfg_attr(rustfmt, rustfmt_skip)]
static LENGTH_TABLE: [u8; 32] = [
    0x0A, 0xFE,
    0x14, 0x02,
    0x28, 0x04,
    0x50, 0x06,
    0xA0, 0x08,
    0x3C, 0x0A,
    0x0E, 0x0C,
    0x1A, 0x0E,
    0x0C, 0x10,
    0x18, 0x12,
    0x30, 0x14,
    0x60, 0x16,
    0xC0, 0x18,
    0x48, 0x1A,
    0x10, 0x1C,
    0x20, 0x1E,
];

///Represents the Length counter used by all NES sound channels except the DMC.
#[derive(Debug)]
pub struct Length {
    halt_bit: usize,
    halted: bool,
    enabled: bool,
    remaining: u8,
}

impl Length {
    pub fn write_halt(&mut self, val: u8) {
        self.halted = (val >> self.halt_bit) & 0x01 != 0;
    }

    pub fn write_counter(&mut self, val: u8) {
        if self.enabled {
            self.remaining = LENGTH_TABLE[(val >> 3) as usize];
        }
    }

    pub fn tick(&mut self) {
        if !self.halted {
            self.remaining = self.remaining.saturating_sub(1);
        }
    }

    pub fn audible(&self) -> bool {
        self.remaining > 0
    }

    pub fn active(&self) -> u8 {
        if self.audible() {
            1
        } else {
            0
        }
    }

    pub fn set_enable(&mut self, enable: bool) {
        self.enabled = enable;
        if !enable {
            self.remaining = 0;
        }
    }

    pub fn new(halt_bit: usize) -> Length {
        Length {
            halt_bit: halt_bit,
            halted: false,
            enabled: false,
            remaining: 0,
        }
    }
}

///Represents the Envelope Generator (volume setting) used by the pulse & noise channels.
#[derive(Debug)]
pub struct Envelope {
    should_loop: bool,
    disable: bool,
    n: u8,

    divider: u8,
    counter: u8,
}

impl Envelope {
    pub fn new() -> Envelope {
        Envelope {
            should_loop: false,
            disable: false,
            n: 0,
            divider: 0,
            counter: 0,
        }
    }

    pub fn write(&mut self, val: u8) {
        self.should_loop = (val >> 5) & 0x01 != 0;
        self.disable = (val >> 6) & 0x01 != 0;
        self.n = val & 0x0F;
        self.divider = self.n;
        self.counter = 15;
    }

    pub fn tick(&mut self) {
        if self.divider == 0 {
            self.envelope_tick();
            self.divider = self.n;
        } else {
            self.divider -= 1;
        }
    }

    pub fn envelope_tick(&mut self) {
        if self.should_loop && self.counter == 0 {
            self.counter = 15;
        } else {
            self.counter = self.counter.saturating_sub(1);
        }
    }

    pub fn volume(&self) -> i16 {
        if self.disable {
            self.n as i16
        } else {
            self.counter as i16
        }
    }
}

#[derive(Debug)]
pub enum TimerClock {
    Clock,
    NoClock,
}

///Represents the CPU-clock timers used by all of the NES channels.
#[derive(Debug)]
pub struct Timer {
    period: u16,
    divider: u32,
    remaining: u32,
}

impl Timer {
    pub fn new(divider: u32) -> Timer {
        Timer {
            period: 0,
            divider: divider,
            remaining: 0,
        }
    }

    pub fn write_low(&mut self, val: u8) {
        self.period = (self.period & 0xFF00) | val as u16;
    }

    pub fn write_high(&mut self, val: u8) {
        self.period = (self.period & 0x00FF) | (val as u16 & 0x0007) << 8;
    }

    pub fn add_period_shift(&mut self, shift: i16) {
        let new_period = (self.period as i16).wrapping_add(shift);
        self.period = new_period as u16;
    }

    pub fn period(&self) -> u16 {
        self.period
    }

    fn wavelen(&self) -> u32 {
        (self.period as u32 + 1) * self.divider
    }

    ///Run the timer until the next clock, or until current_cyc reaches to_cycle.
    ///Returns either Clock or NoClock depending on if it reached a clock or not.
    pub fn run(&mut self, current_cyc: &mut u32, to_cyc: u32) -> TimerClock {
        let end_wavelen = *current_cyc + self.remaining;

        if end_wavelen < to_cyc {
            *current_cyc += self.remaining;
            self.remaining = self.wavelen();
            TimerClock::Clock
        } else {
            self.remaining -= to_cyc - *current_cyc;
            *current_cyc = to_cyc;
            TimerClock::NoClock
        }
    }
}
<commit_msg>Fix music bug in super mario bros.<commit_after>//! This module contains implementations of the common components used by the
//! various NES sound channels.

#[cfg_attr(rustfmt, rustfmt_skip)]
static LENGTH_TABLE: [u8; 32] = [
    0x0A, 0xFE,
    0x14, 0x02,
    0x28, 0x04,
    0x50, 0x06,
    0xA0, 0x08,
    0x3C, 0x0A,
    0x0E, 0x0C,
    0x1A, 0x0E,
    0x0C, 0x10,
    0x18, 0x12,
    0x30, 0x14,
    0x60, 0x16,
    0xC0, 0x18,
    0x48, 0x1A,
    0x10, 0x1C,
    0x20, 0x1E,
];

///Represents the Length counter used by all NES sound channels except the DMC.
#[derive(Debug)]
pub struct Length {
    halt_bit: usize,
    halted: bool,
    enabled: bool,
    remaining: u8,
}

impl Length {
    pub fn write_halt(&mut self, val: u8) {
        self.halted = (val >> self.halt_bit) & 0x01 != 0;
    }

    pub fn write_counter(&mut self, val: u8) {
        if self.enabled {
            self.remaining = LENGTH_TABLE[(val >> 3) as usize];
        }
    }

    pub fn tick(&mut self) {
        if !self.halted {
            self.remaining = self.remaining.saturating_sub(1);
        }
    }

    pub fn audible(&self) -> bool {
        self.remaining > 0
    }

    pub fn active(&self) -> u8 {
        if self.audible() {
            1
        } else {
            0
        }
    }

    pub fn set_enable(&mut self, enable: bool) {
        self.enabled = enable;
        if !enable {
            self.remaining = 0;
        }
    }

    pub fn new(halt_bit: usize) -> Length {
        Length {
            halt_bit: halt_bit,
            halted: false,
            enabled: false,
            remaining: 0,
        }
    }
}

///Represents the Envelope Generator (volume setting) used by the pulse & noise channels.
#[derive(Debug)]
pub struct Envelope {
    should_loop: bool,
    constant_volume: bool,
    n: u8,

    divider: u8,
    counter: u8,
}

impl Envelope {
    pub fn new() -> Envelope {
        Envelope {
            should_loop: false,
            constant_volume: false,
            n: 0,
            divider: 0,
            counter: 0,
        }
    }

    pub fn write(&mut self, val: u8) {
        self.should_loop = (val >> 5) & 0x01 != 0;
        self.constant_volume = (val >> 4) & 0x01 != 0;
        self.n = val & 0x0F;
        self.divider = self.n;
        self.counter = 15;
    }

    pub fn tick(&mut self) {
        if self.divider == 0 {
            self.envelope_tick();
            self.divider = self.n;
        } else {
            self.divider -= 1;
        }
    }

    pub fn envelope_tick(&mut self) {
        if self.should_loop && self.counter == 0 {
            self.counter = 15;
        } else {
            self.counter = self.counter.saturating_sub(1);
        }
    }

    pub fn volume(&self) -> i16 {
        if self.constant_volume {
            self.n as i16
        } else {
            self.counter as i16
        }
    }
}

#[derive(Debug)]
pub enum TimerClock {
    Clock,
    NoClock,
}

///Represents the CPU-clock timers used by all of the NES channels.
#[derive(Debug)]
pub struct Timer {
    period: u16,
    divider: u32,
    remaining: u32,
}

impl Timer {
    pub fn new(divider: u32) -> Timer {
        Timer {
            period: 0,
            divider: divider,
            remaining: 0,
        }
    }

    pub fn write_low(&mut self, val: u8) {
        self.period = (self.period & 0xFF00) | val as u16;
    }

    pub fn write_high(&mut self, val: u8) {
        self.period = (self.period & 0x00FF) | (val as u16 & 0x0007) << 8;
    }

    pub fn add_period_shift(&mut self, shift: i16) {
        let new_period = (self.period as i16).wrapping_add(shift);
        self.period = new_period as u16;
    }

    pub fn period(&self) -> u16 {
        self.period
    }

    fn wavelen(&self) -> u32 {
        (self.period as u32 + 1) * self.divider
    }

    ///Run the timer until the next clock, or until current_cyc reaches to_cycle.
    ///Returns either Clock or NoClock depending on if it reached a clock or not.
    pub fn run(&mut self, current_cyc: &mut u32, to_cyc: u32) -> TimerClock {
        let end_wavelen = *current_cyc + self.remaining;

        if end_wavelen < to_cyc {
            *current_cyc += self.remaining;
            self.remaining = self.wavelen();
            TimerClock::Clock
        } else {
            self.remaining -= to_cyc - *current_cyc;
            *current_cyc = to_cyc;
            TimerClock::NoClock
        }
    }
}
<|endoftext|>


<commit_before>use std::ffi::OsString;
use std::path::Path;

use log::warn;

pub struct ExtAttr {
    pub key: OsString,
    pub value: Option<Vec<u8>>,
}

pub fn ext_attrs<'p, P>(path: &'p P) -> std::io::Result<ExtAttrs<'p>>
where
    P: AsRef<Path>,
{
    let iter = xattr::list(&path)?;

    Ok(ExtAttrs {
        path: path.as_ref(),
        iter: iter,
    })
}

pub struct ExtAttrs<'p> {
    path: &'p Path,
    iter: xattr::XAttrs,
}

impl<'p> Iterator for ExtAttrs<'p> {

    type Item = ExtAttr;

    fn next(&mut self) -> Option<ExtAttr> {
        for key in &mut self.iter {
            let value = match xattr::get(self.path, &key) {
                Ok(value) => value,
                Err(error) => {
                    warn! {
                        "failed to collect {key:?} of '{path}': {cause}",
                        key = key,
                        path = self.path.display(),
                        cause = error,
                    };
                    continue
                },
            };

            return Some(ExtAttr {
                key: key,
                value: value,
            });
        }

        None
    }
}
<commit_msg>Beautify the extended attribute collection code.<commit_after>use std::ffi::{OsStr, OsString};
use std::path::Path;

use log::warn;

pub struct ExtAttr {
    pub key: OsString,
    pub value: Option<Vec<u8>>,
}

pub fn ext_attrs<'p, P>(path: &'p P) -> std::io::Result<ExtAttrs<'p>>
where
    P: AsRef<Path>,
{
    let iter = xattr::list(&path)?;

    Ok(ExtAttrs {
        path: path.as_ref(),
        iter: iter,
    })
}

pub struct ExtAttrs<'p> {
    path: &'p Path,
    iter: xattr::XAttrs,
}

impl<'p> Iterator for ExtAttrs<'p> {

    type Item = ExtAttr;

    fn next(&mut self) -> Option<ExtAttr> {
        for key in &mut self.iter {
            let value = match ext_attr_value(self.path, &key) {
                Ok(value) => value,
                Err(()) => continue,
            };

            return Some(ExtAttr {
                key: key,
                value: value,
            });
        }

        None
    }
}

fn ext_attr_value<P>(path: P, key: &OsStr) -> Result<Option<Vec<u8>>, ()>
where
    P: AsRef<Path>,
{
    xattr::get(&path, key).map_err(|error| {
        warn! {
            "failed to collect {key:?} of '{path}': {cause}",
            key = key,
            path = path.as_ref().display(),
            cause = error,
        };
    })
}
<|endoftext|>


<commit_before>//! Define trait for Hermite matrices

use ndarray::{Ix2, Array, LinalgScalar};
use std::fmt::Debug;
use num_traits::float::Float;
use lapack::c::Layout;

use matrix::Matrix;
use square::SquareMatrix;
use error::LinalgError;
use eigh::ImplEigh;
use qr::ImplQR;
use svd::ImplSVD;
use norm::ImplNorm;
use solve::ImplSolve;
use cholesky::ImplCholesky;

/// Methods for Hermite matrix
pub trait HermiteMatrix: SquareMatrix + Matrix {
    /// eigenvalue decomposition
    fn eigh(self) -> Result<(Self::Vector, Self), LinalgError>;
    /// symmetric square root of Hermite matrix
    fn ssqrt(self) -> Result<Self, LinalgError>;
    /// Cholesky factorization
    fn cholesky(self) -> Result<Self, LinalgError>;
}

impl<A> HermiteMatrix for Array<A, Ix2>
    where A: ImplQR + ImplSVD + ImplNorm + ImplSolve + ImplEigh + ImplCholesky + LinalgScalar + Float + Debug
{
    fn eigh(self) -> Result<(Self::Vector, Self), LinalgError> {
        self.check_square()?;
        let layout = self.layout()?;
        let (rows, cols) = self.size();
        let (w, a) = ImplEigh::eigh(layout, rows, self.into_raw_vec())?;
        let ea = Array::from_vec(w);
        let va = match layout {
            Layout::ColumnMajor => Array::from_vec(a).into_shape((rows, cols)).unwrap().reversed_axes(),
            Layout::RowMajor => Array::from_vec(a).into_shape((rows, cols)).unwrap(),
        };
        Ok((ea, va))
    }
    fn ssqrt(self) -> Result<Self, LinalgError> {
        let (n, _) = self.size();
        let (e, v) = self.eigh()?;
        let mut res = Array::zeros((n, n));
        for i in 0..n {
            for j in 0..n {
                res[(i, j)] = e[i].sqrt() * v[(j, i)];
            }
        }
        Ok(v.dot(&res))
    }
    fn cholesky(self) -> Result<Self, LinalgError> {
        self.check_square()?;
        let (n, _) = self.size();
        let layout = self.layout()?;
        let a = ImplCholesky::cholesky(layout, n, self.into_raw_vec())?;
        let mut c = match layout {
            Layout::RowMajor => Array::from_vec(a).into_shape((n, n)).unwrap(),
            Layout::ColumnMajor => Array::from_vec(a).into_shape((n, n)).unwrap().reversed_axes(),
        };
        for ((i, j), val) in c.indexed_iter_mut() {
            if i > j {
                *val = A::zero();
            }
        }
        Ok(c)
    }
}
<commit_msg>Implement determinant for Hermite<commit_after>//! Define trait for Hermite matrices

use ndarray::{Ix2, Array, LinalgScalar};
use std::fmt::Debug;
use num_traits::float::Float;
use num_traits::One;
use lapack::c::Layout;

use matrix::Matrix;
use square::SquareMatrix;
use error::LinalgError;
use eigh::ImplEigh;
use qr::ImplQR;
use svd::ImplSVD;
use norm::ImplNorm;
use solve::ImplSolve;
use cholesky::ImplCholesky;

/// Methods for Hermite matrix
pub trait HermiteMatrix: SquareMatrix + Matrix {
    /// eigenvalue decomposition
    fn eigh(self) -> Result<(Self::Vector, Self), LinalgError>;
    /// symmetric square root of Hermite matrix
    fn ssqrt(self) -> Result<Self, LinalgError>;
    /// Cholesky factorization
    fn cholesky(self) -> Result<Self, LinalgError>;
    /// calc determinant using Cholesky factorization
    fn deth(self) -> Result<Self::Scalar, LinalgError>;
}

impl<A> HermiteMatrix for Array<A, Ix2>
    where A: ImplQR + ImplSVD + ImplNorm + ImplSolve + ImplEigh + ImplCholesky + LinalgScalar + Float + Debug + One
{
    fn eigh(self) -> Result<(Self::Vector, Self), LinalgError> {
        self.check_square()?;
        let layout = self.layout()?;
        let (rows, cols) = self.size();
        let (w, a) = ImplEigh::eigh(layout, rows, self.into_raw_vec())?;
        let ea = Array::from_vec(w);
        let va = match layout {
            Layout::ColumnMajor => Array::from_vec(a).into_shape((rows, cols)).unwrap().reversed_axes(),
            Layout::RowMajor => Array::from_vec(a).into_shape((rows, cols)).unwrap(),
        };
        Ok((ea, va))
    }
    fn ssqrt(self) -> Result<Self, LinalgError> {
        let (n, _) = self.size();
        let (e, v) = self.eigh()?;
        let mut res = Array::zeros((n, n));
        for i in 0..n {
            for j in 0..n {
                res[(i, j)] = e[i].sqrt() * v[(j, i)];
            }
        }
        Ok(v.dot(&res))
    }
    fn cholesky(self) -> Result<Self, LinalgError> {
        self.check_square()?;
        let (n, _) = self.size();
        let layout = self.layout()?;
        let a = ImplCholesky::cholesky(layout, n, self.into_raw_vec())?;
        let mut c = match layout {
            Layout::RowMajor => Array::from_vec(a).into_shape((n, n)).unwrap(),
            Layout::ColumnMajor => Array::from_vec(a).into_shape((n, n)).unwrap().reversed_axes(),
        };
        for ((i, j), val) in c.indexed_iter_mut() {
            if i > j {
                *val = A::zero();
            }
        }
        Ok(c)
    }
    fn deth(self) -> Result<Self::Scalar, LinalgError> {
        let (n, _) = self.size();
        let c = self.cholesky()?;
        let rt = (0..n).map(|i| c[(i, i)]).fold(A::one(), |det, c| det * c);
        Ok(rt * rt)
    }
}
<|endoftext|>


<commit_before>//! Reimplements core logic and types from `std::io` in an `alloc`-friendly
//! fashion.

use lib::*;

pub enum ErrorKind {
    Other,
}

pub struct Error(&'static str);

impl Display for Error {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        Display::fmt(self.0, formatter)
    }
}

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

impl Error {
    pub(crate) fn new(kind: ErrorKind, error: &'static str) -> Error {
        let _ = kind;
        Error(error)
    }
}

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

pub trait Write {
    fn write(&mut self, buf: &[u8]) -> Result<usize>;

    fn write_all(&mut self, buf: &[u8]) -> Result<()> {
        // All our Write impls in no_std mode always write the whole buffer in
        // one call infallibly.
        let result = self.write(buf);
        debug_assert!(result.is_ok());
        debug_assert_eq!(result.unwrap(), buf.len());
        Ok(())
    }

    fn flush(&mut self) -> Result<()>;
}

impl<W: Write> Write for &mut W {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        (*self).write(buf)
    }

    #[inline]
    fn write_all(&mut self, buf: &[u8]) -> Result<()> {
        (*self).write_all(buf)
    }

    #[inline]
    fn flush(&mut self) -> Result<()> {
        (*self).flush()
    }
}

impl Write for Vec<u8> {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        self.extend_from_slice(buf);
        Ok(buf.len())
    }

    #[inline]
    fn write_all(&mut self, buf: &[u8]) -> Result<()> {
        self.extend_from_slice(buf);
        Ok(())
    }

    #[inline]
    fn flush(&mut self) -> Result<()> {
        Ok(())
    }
}

pub trait Read {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize>;
    fn bytes(self) -> Bytes<Self>
    where
        Self: Sized,
    {
        Bytes { inner: self }
    }
}

pub struct Bytes<R> {
    inner: R,
}

impl<R: Read> Iterator for Bytes<R> {
    type Item = Result<u8>;

    fn next(&mut self) -> Option<Result<u8>> {
        let mut byte = 0;
        match self.inner.read(slice::from_mut(&mut byte)) {
            Ok(0) => None,
            Ok(..) => Some(Ok(byte)),
            Err(e) => Some(Err(e)),
        }
    }
}
<commit_msg>Remove unneeded Debug impl on io error<commit_after>//! Reimplements core logic and types from `std::io` in an `alloc`-friendly
//! fashion.

use lib::*;

pub enum ErrorKind {
    Other,
}

pub struct Error(&'static str);

impl Display for Error {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        Display::fmt(self.0, formatter)
    }
}

impl Error {
    pub(crate) fn new(kind: ErrorKind, error: &'static str) -> Error {
        let _ = kind;
        Error(error)
    }
}

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

pub trait Write {
    fn write(&mut self, buf: &[u8]) -> Result<usize>;

    fn write_all(&mut self, buf: &[u8]) -> Result<()> {
        // All our Write impls in no_std mode always write the whole buffer in
        // one call infallibly.
        let result = self.write(buf);
        debug_assert!(result.is_ok());
        debug_assert_eq!(result.unwrap_or(0), buf.len());
        Ok(())
    }

    fn flush(&mut self) -> Result<()>;
}

impl<W: Write> Write for &mut W {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        (*self).write(buf)
    }

    #[inline]
    fn write_all(&mut self, buf: &[u8]) -> Result<()> {
        (*self).write_all(buf)
    }

    #[inline]
    fn flush(&mut self) -> Result<()> {
        (*self).flush()
    }
}

impl Write for Vec<u8> {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        self.extend_from_slice(buf);
        Ok(buf.len())
    }

    #[inline]
    fn write_all(&mut self, buf: &[u8]) -> Result<()> {
        self.extend_from_slice(buf);
        Ok(())
    }

    #[inline]
    fn flush(&mut self) -> Result<()> {
        Ok(())
    }
}

pub trait Read {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize>;
    fn bytes(self) -> Bytes<Self>
    where
        Self: Sized,
    {
        Bytes { inner: self }
    }
}

pub struct Bytes<R> {
    inner: R,
}

impl<R: Read> Iterator for Bytes<R> {
    type Item = Result<u8>;

    fn next(&mut self) -> Option<Result<u8>> {
        let mut byte = 0;
        match self.inner.read(slice::from_mut(&mut byte)) {
            Ok(0) => None,
            Ok(..) => Some(Ok(byte)),
            Err(e) => Some(Err(e)),
        }
    }
}
<|endoftext|>


<commit_before>/// Implement +, -, *, / operators for Tensor
/// +=, -= are provided as methods of ops::inplace_add, ops::inplace_sub
extern crate ndarray;

use ndarray_ext::NdArray;
use ops;
use std::mem;
use std::ops::{Add, Div, Mul, Sub};
use std::result::Result;
use tensor::Tensor;

pub struct AddOp;
pub struct SubOp;
pub struct MulOp;
pub struct DivOp;
pub struct InplaceAddOp;
pub struct InplaceSubOp;

impl ops::Op for AddOp {
    fn name(&self) -> &str
    {
        "Add"
    }

    fn compute(&self, xs: &[&NdArray]) -> Result<NdArray, ::OpComputeErrorStatus>
    {
        add_forward(xs[0], xs[1])
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let (gy1, gy2) = maybe_reduce_gy(inputs[0], inputs[1], gy);
        vec![Some(gy1), Some(gy2)]
    }
}

impl ops::Op for SubOp {
    fn name(&self) -> &str
    {
        "Sub"
    }

    fn compute(&self, xs: &[&NdArray]) -> Result<NdArray, ::OpComputeErrorStatus>
    {
        let x0 = xs[0];
        let x1 = xs[1];
        let shape0: &[usize] = x0.shape();
        if shape0 == &[] {
            // a is scalar
            let x0_elem = x0[ndarray::IxDyn(&[])];
            Ok(x1.map(move |a| x0_elem - a))
        } else {
            Ok(x0 - x1)
        }
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let (gy1, gy2) = maybe_reduce_gy(inputs[0], inputs[1], gy);
        vec![Some(gy1), Some(gy2 * -1.)]
    }
}

impl ops::Op for MulOp {
    fn name(&self) -> &str
    {
        "Mul"
    }

    fn compute(&self, xs: &[&NdArray]) -> Result<NdArray, ::OpComputeErrorStatus>
    {
        mul_forward(xs[0], xs[1])
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let x0 = inputs[0];
        let x1 = inputs[1];
        let (gy1, gy2) = maybe_reduce_gy(x0, x1, gy);
        vec![Some(gy1 * x1), Some(gy2 * x0)]
    }
}

impl ops::Op for DivOp {
    fn name(&self) -> &str
    {
        "Div"
    }

    fn compute(&self, xs: &[&NdArray]) -> Result<NdArray, ::OpComputeErrorStatus>
    {
        let x0 = xs[0];
        let x1 = xs[1];
        let shape0: &[usize] = x0.shape();
        if shape0 == &[] {
            // a is scalar
            let x0_elem = x0[ndarray::IxDyn(&[])];
            Ok(x1.map(move |a| x0_elem / a))
        } else {
            Ok(x0 / x1)
        }
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let x0 = inputs[0];
        let x1 = inputs[1];
        let (gy1, gy2) = maybe_reduce_gy(x0, x1, gy);
        vec![Some(gy1 / x1), Some(-1 * x0 * ops::pow(x1, -2.) * gy2)]
    }
}

impl ops::Op for InplaceAddOp {
    fn inplace(&self) -> bool
    {
        true
    }

    fn name(&self) -> &str
    {
        "InplaceAddOp"
    }

    fn compute_inplace(&self, xs: &mut [&mut NdArray]) -> Result<(), ::OpComputeErrorStatus>
    {
        // safe transmute probably
        let x1: &&NdArray = unsafe { mem::transmute(&mut xs[1]) };
        let x0 = &mut xs[0];
        x0.zip_mut_with(x1, |a, &b| *a += b);
        Ok(())
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let (gy1, gy2) = maybe_reduce_gy(inputs[0], inputs[1], gy);
        vec![Some(gy1), Some(gy2)]
    }
}

impl ops::Op for InplaceSubOp {
    fn inplace(&self) -> bool
    {
        true
    }

    fn name(&self) -> &str
    {
        "InplaceSubOp"
    }

    fn compute_inplace(&self, xs: &mut [&mut NdArray]) -> Result<(), ::OpComputeErrorStatus>
    {
        // safe transmute probably
        let x1: &&NdArray = unsafe { mem::transmute(&mut xs[1]) };
        let x0 = &mut xs[0];
        x0.zip_mut_with(x1, |a, &b| *a -= b);
        Ok(())
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let (gy1, gy2) = maybe_reduce_gy(inputs[0], inputs[1], gy);
        vec![Some(gy1), Some(-1. * gy2)]
    }
}

#[inline]
// Reduce gy if broadcast occurred in the forward pass
fn maybe_reduce_gy(x0: &Tensor, x1: &Tensor, gy: &Tensor) -> (Tensor, Tensor)
{
    let shape0 = x0.shape();
    let shape1 = x1.shape();
    let gy_shape = gy.shape();
    let sum0 = ops::reduction_ops::ReduceSum { keep_dims: true, sparse_axes: true };
    let sum1 = ops::reduction_ops::ReduceSum { keep_dims: true, sparse_axes: true };
    return (
        ops::apply_op(
            sum0,
            &[gy, &ops::not_equal(&gy_shape, &shape0)],
            Some(shape0),
        ), // gy1
        ops::apply_op(
            sum1,
            &[gy, &ops::not_equal(&gy_shape, &shape1)],
            Some(shape1),
        ), /* gy2 */
    );
}

macro_rules! impl_scalar_to_tensor {
    ($func_name:ident, $scalar_type:ty) => {
        fn $func_name(val: $scalar_type) -> Tensor
        {
            ops::scalar(val as f32)
        }
    };
}

// -- std::ops::{Add, Sub, Mul, Div} implementations --

macro_rules! impl_bin_op_between_tensor_and_scalar {
    (
        $trt:ident,
        $func:ident,
        $op:ident,
        $scalar_type:ty,
        $conversion_fn:ident
    ) => {

        // scalar op Tensor
        impl $trt<Tensor> for $scalar_type {
            type Output = Tensor;
            fn $func(self, rhs: Tensor) -> Self::Output
            {
                ops::apply_op($op, &[&$conversion_fn(self), &rhs], Some(rhs.shape()))
            }
        }

        // scalar op &Tensor
        impl<'a> $trt<&'a Tensor> for $scalar_type {
            type Output = Tensor;
            fn $func(self, rhs: &'a Tensor) -> Self::Output
            {
                ops::apply_op($op, &[&$conversion_fn(self), rhs], Some(rhs.shape()))
            }
        }

        // Tensor op scalar
        impl $trt<$scalar_type> for Tensor {
            type Output = Tensor;
            fn $func(self, rhs: $scalar_type) -> Self::Output
            {
                ops::apply_op($op, &[&self, &$conversion_fn(rhs)], Some(self.shape()))
            }
        }

        // &Tensor op scalar
        impl<'a> $trt<$scalar_type> for &'a Tensor {
            type Output = Tensor;
            fn $func(self, rhs: $scalar_type) -> Self::Output
            {
                ops::apply_op($op, &[self, &$conversion_fn(rhs)], Some(self.shape()))
            }
        }
    }
}

macro_rules! impl_bin_op_between_tensors {
    (
        $trt:ident,
        $func:ident,
        $op:ident
    ) => {
        // Tensor op Tensor
        impl $trt for Tensor {
            type Output = Tensor;
            fn $func(self, rhs: Tensor) -> Self::Output
            {
                ops::$func(&self, &rhs)
            }
        }

        // Tensor op &Tensor
        impl<'a> $trt<&'a Tensor> for Tensor {
            type Output = Tensor;
            fn $func(self, rhs: &Tensor) -> Self::Output
            {
                ops::$func(&self, rhs)
            }
        }

        // &Tensor op Tensor
        impl<'a> $trt<Tensor> for &'a Tensor {
            type Output = Tensor;
            fn $func(self, rhs: Tensor) -> Self::Output
            {
                ops::$func(&self, &rhs)
            }
        }

        // &Tensor op &Tensor
        // lifetime of the two tensors are unrelated
        impl<'a, 'b> $trt<&'a Tensor> for &'b Tensor {
            type Output = Tensor;
            fn $func(self, rhs: &Tensor) -> Self::Output
            {
                ops::$func(self, rhs)
            }
        }
    };
}

macro_rules! impl_bin_op_forward {
    ($forward_name:ident, $bin_op:tt) => {
        fn $forward_name(x0: &NdArray, x1: &NdArray) -> Result<NdArray, ::OpComputeErrorStatus>
        {
            let shape0: &[usize]  = x0.shape();
            let shape1: &[usize]  = x1.shape();
            let len0: usize = shape0.iter().product();
            let len1: usize = shape1.iter().product();
            let (larger, smaller) = if len0 > len1 { (x0, x1) } else { (x1, x0) };
            let smaller_sh = if len0 > len1 { shape1 } else { shape0 };

            if smaller_sh == &[] {
                let smaller_elem = smaller[ndarray::IxDyn(&[])];
                Ok(larger.map(move |a| smaller_elem $bin_op a))
            } else {
                Ok(larger $bin_op smaller)
            }
        }
    };
}

impl_bin_op_forward!(add_forward, +);
impl_bin_op_forward!(mul_forward, *);

impl_scalar_to_tensor!(isize_to_tensor, isize);
impl_scalar_to_tensor!(usize_to_tensor, usize);
impl_scalar_to_tensor!(i32_to_tensor, i32);
impl_scalar_to_tensor!(i64_to_tensor, i64);
impl_scalar_to_tensor!(u32_to_tensor, u32);
impl_scalar_to_tensor!(u64_to_tensor, u64);
impl_scalar_to_tensor!(f32_to_tensor, f32);
impl_scalar_to_tensor!(f64_to_tensor, f64);

impl_bin_op_between_tensors!(Add, add, AddOp);
impl_bin_op_between_tensors!(Sub, sub, SubOp);
impl_bin_op_between_tensors!(Mul, mul, MulOp);
impl_bin_op_between_tensors!(Div, div, DivOp);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, i32, i32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, i32, i32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, i32, i32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, i32, i32_to_tensor);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, i64, i64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, i64, i64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, i64, i64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, i64, i64_to_tensor);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, f32, f32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, f32, f32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, f32, f32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, f32, f32_to_tensor);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, f64, f64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, f64, f64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, f64, f64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, f64, f64_to_tensor);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, u32, u32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, u32, u32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, u32, u32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, u32, u32_to_tensor);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, u64, u64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, u64, u64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, u64, u64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, u64, u64_to_tensor);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, usize, usize_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, usize, usize_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, usize, usize_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, usize, usize_to_tensor);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, isize, isize_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, isize, isize_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, isize, isize_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, isize, isize_to_tensor);
<commit_msg>Remove redundant macro<commit_after>/// Implement +, -, *, / operators for Tensor
/// +=, -= are provided as methods of ops::inplace_add, ops::inplace_sub
extern crate ndarray;

use ndarray_ext::NdArray;
use ops;
use std::mem;
use std::ops::{Add, Div, Mul, Sub};
use std::result::Result;
use tensor::Tensor;

pub struct AddOp;
pub struct SubOp;
pub struct MulOp;
pub struct DivOp;
pub struct InplaceAddOp;
pub struct InplaceSubOp;

impl ops::Op for AddOp {
    fn name(&self) -> &str
    {
        "Add"
    }

    fn compute(&self, xs: &[&NdArray]) -> Result<NdArray, ::OpComputeErrorStatus>
    {
        add_forward(xs[0], xs[1])
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let (gy1, gy2) = maybe_reduce_gy(inputs[0], inputs[1], gy);
        vec![Some(gy1), Some(gy2)]
    }
}

impl ops::Op for SubOp {
    fn name(&self) -> &str
    {
        "Sub"
    }

    fn compute(&self, xs: &[&NdArray]) -> Result<NdArray, ::OpComputeErrorStatus>
    {
        let x0 = xs[0];
        let x1 = xs[1];
        let shape0: &[usize] = x0.shape();
        if shape0 == &[] {
            // a is scalar
            let x0_elem = x0[ndarray::IxDyn(&[])];
            Ok(x1.map(move |a| x0_elem - a))
        } else {
            Ok(x0 - x1)
        }
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let (gy1, gy2) = maybe_reduce_gy(inputs[0], inputs[1], gy);
        vec![Some(gy1), Some(gy2 * -1.)]
    }
}

impl ops::Op for MulOp {
    fn name(&self) -> &str
    {
        "Mul"
    }

    fn compute(&self, xs: &[&NdArray]) -> Result<NdArray, ::OpComputeErrorStatus>
    {
        mul_forward(xs[0], xs[1])
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let x0 = inputs[0];
        let x1 = inputs[1];
        let (gy1, gy2) = maybe_reduce_gy(x0, x1, gy);
        vec![Some(gy1 * x1), Some(gy2 * x0)]
    }
}

impl ops::Op for DivOp {
    fn name(&self) -> &str
    {
        "Div"
    }

    fn compute(&self, xs: &[&NdArray]) -> Result<NdArray, ::OpComputeErrorStatus>
    {
        let x0 = xs[0];
        let x1 = xs[1];
        let shape0: &[usize] = x0.shape();
        if shape0 == &[] {
            // a is scalar
            let x0_elem = x0[ndarray::IxDyn(&[])];
            Ok(x1.map(move |a| x0_elem / a))
        } else {
            Ok(x0 / x1)
        }
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let x0 = inputs[0];
        let x1 = inputs[1];
        let (gy1, gy2) = maybe_reduce_gy(x0, x1, gy);
        vec![Some(gy1 / x1), Some(-1 * x0 * ops::pow(x1, -2.) * gy2)]
    }
}

impl ops::Op for InplaceAddOp {
    fn inplace(&self) -> bool
    {
        true
    }

    fn name(&self) -> &str
    {
        "InplaceAddOp"
    }

    fn compute_inplace(&self, xs: &mut [&mut NdArray]) -> Result<(), ::OpComputeErrorStatus>
    {
        // safe transmute probably
        let x1: &&NdArray = unsafe { mem::transmute(&mut xs[1]) };
        let x0 = &mut xs[0];
        x0.zip_mut_with(x1, |a, &b| *a += b);
        Ok(())
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let (gy1, gy2) = maybe_reduce_gy(inputs[0], inputs[1], gy);
        vec![Some(gy1), Some(gy2)]
    }
}

impl ops::Op for InplaceSubOp {
    fn inplace(&self) -> bool
    {
        true
    }

    fn name(&self) -> &str
    {
        "InplaceSubOp"
    }

    fn compute_inplace(&self, xs: &mut [&mut NdArray]) -> Result<(), ::OpComputeErrorStatus>
    {
        // safe transmute probably
        let x1: &&NdArray = unsafe { mem::transmute(&mut xs[1]) };
        let x0 = &mut xs[0];
        x0.zip_mut_with(x1, |a, &b| *a -= b);
        Ok(())
    }

    fn grad(&self, gy: &Tensor, inputs: &[&Tensor], _: &Tensor) -> Vec<Option<Tensor>>
    {
        let (gy1, gy2) = maybe_reduce_gy(inputs[0], inputs[1], gy);
        vec![Some(gy1), Some(-1. * gy2)]
    }
}

#[inline]
// Reduce gy if broadcast occurred in the forward pass
fn maybe_reduce_gy(x0: &Tensor, x1: &Tensor, gy: &Tensor) -> (Tensor, Tensor)
{
    let shape0 = x0.shape();
    let shape1 = x1.shape();
    let gy_shape = gy.shape();
    let sum0 = ops::reduction_ops::ReduceSum { keep_dims: true, sparse_axes: true };
    let sum1 = ops::reduction_ops::ReduceSum { keep_dims: true, sparse_axes: true };
    return (
        ops::apply_op(
            sum0,
            &[gy, &ops::not_equal(&gy_shape, &shape0)],
            Some(shape0),
        ), // gy1
        ops::apply_op(
            sum1,
            &[gy, &ops::not_equal(&gy_shape, &shape1)],
            Some(shape1),
        ), /* gy2 */
    );
}

// -- std::ops::{Add, Sub, Mul, Div} implementations --

macro_rules! impl_bin_op_between_tensor_and_scalar {
    (
        $trt:ident,
        $func:ident,
        $op:ident,
        $scalar_type:ty,
        $conversion_fn:ident
    ) => {

        // scalar op Tensor
        impl $trt<Tensor> for $scalar_type {
            type Output = Tensor;
            fn $func(self, rhs: Tensor) -> Self::Output
            {
                ops::apply_op($op, &[&ops::scalar(self as f32), &rhs], Some(rhs.shape()))
            }
        }

        // scalar op &Tensor
        impl<'a> $trt<&'a Tensor> for $scalar_type {
            type Output = Tensor;
            fn $func(self, rhs: &'a Tensor) -> Self::Output
            {
                ops::apply_op($op, &[&ops::scalar(self as f32), rhs], Some(rhs.shape()))
            }
        }

        // Tensor op scalar
        impl $trt<$scalar_type> for Tensor {
            type Output = Tensor;
            fn $func(self, rhs: $scalar_type) -> Self::Output
            {
                ops::apply_op($op, &[&self, &ops::scalar(rhs as f32)], Some(self.shape()))
            }
        }

        // &Tensor op scalar
        impl<'a> $trt<$scalar_type> for &'a Tensor {
            type Output = Tensor;
            fn $func(self, rhs: $scalar_type) -> Self::Output
            {
                ops::apply_op($op, &[self, &ops::scalar(rhs as f32)], Some(self.shape()))
            }
        }
    }
}

macro_rules! impl_bin_op_between_tensors {
    (
        $trt:ident,
        $func:ident,
        $op:ident
    ) => {
        // Tensor op Tensor
        impl $trt for Tensor {
            type Output = Tensor;
            fn $func(self, rhs: Tensor) -> Self::Output
            {
                ops::$func(&self, &rhs)
            }
        }

        // Tensor op &Tensor
        impl<'a> $trt<&'a Tensor> for Tensor {
            type Output = Tensor;
            fn $func(self, rhs: &Tensor) -> Self::Output
            {
                ops::$func(&self, rhs)
            }
        }

        // &Tensor op Tensor
        impl<'a> $trt<Tensor> for &'a Tensor {
            type Output = Tensor;
            fn $func(self, rhs: Tensor) -> Self::Output
            {
                ops::$func(&self, &rhs)
            }
        }

        // &Tensor op &Tensor
        // lifetime of the two tensors are unrelated
        impl<'a, 'b> $trt<&'a Tensor> for &'b Tensor {
            type Output = Tensor;
            fn $func(self, rhs: &Tensor) -> Self::Output
            {
                ops::$func(self, rhs)
            }
        }
    };
}

macro_rules! impl_bin_op_forward {
    ($forward_name:ident, $bin_op:tt) => {
        fn $forward_name(x0: &NdArray, x1: &NdArray) -> Result<NdArray, ::OpComputeErrorStatus>
        {
            let shape0: &[usize]  = x0.shape();
            let shape1: &[usize]  = x1.shape();
            let len0: usize = shape0.iter().product();
            let len1: usize = shape1.iter().product();
            let (larger, smaller) = if len0 > len1 { (x0, x1) } else { (x1, x0) };
            let smaller_sh = if len0 > len1 { shape1 } else { shape0 };

            if smaller_sh == &[] {
                let smaller_elem = smaller[ndarray::IxDyn(&[])];
                Ok(larger.map(move |a| smaller_elem $bin_op a))
            } else {
                Ok(larger $bin_op smaller)
            }
        }
    };
}

impl_bin_op_forward!(add_forward, +);
impl_bin_op_forward!(mul_forward, *);

impl_bin_op_between_tensors!(Add, add, AddOp);
impl_bin_op_between_tensors!(Sub, sub, SubOp);
impl_bin_op_between_tensors!(Mul, mul, MulOp);
impl_bin_op_between_tensors!(Div, div, DivOp);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, i32, i32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, i32, i32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, i32, i32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, i32, i32_to_tensor);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, i64, i64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, i64, i64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, i64, i64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, i64, i64_to_tensor);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, f32, f32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, f32, f32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, f32, f32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, f32, f32_to_tensor);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, f64, f64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, f64, f64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, f64, f64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, f64, f64_to_tensor);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, u32, u32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, u32, u32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, u32, u32_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, u32, u32_to_tensor);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, u64, u64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, u64, u64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, u64, u64_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, u64, u64_to_tensor);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, usize, usize_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, usize, usize_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, usize, usize_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, usize, usize_to_tensor);

impl_bin_op_between_tensor_and_scalar!(Add, add, AddOp, isize, isize_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Sub, sub, SubOp, isize, isize_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Mul, mul, MulOp, isize, isize_to_tensor);
impl_bin_op_between_tensor_and_scalar!(Div, div, DivOp, isize, isize_to_tensor);
<|endoftext|>


<commit_before>
use std::ops::Fn;
use regex::{Captures, Regex};

/// A Parser is designed to take an input type and turn part of it into an
/// output type.  Probably in every real-life scenario, the input type is a
/// slice of some type.  Thus a parser returns an output type along with another
/// input type, which for slices would be the rest of the data.
pub trait Parser<'a, I, O> {

  fn parse(&self, data: I) -> ParseResult<'a, I, O>;

}

pub type ParseResult<'a, I:'a, O> = Result<(O, I), String>;

pub fn literal<'a, T:'a + Eq + Clone>(literal: T) -> LiteralParser<'a, T>{
  LiteralParser{literal: literal}
}

/// A parser that consumes one item in a slice and only returns ok if it equals
/// the given literal
pub struct LiteralParser<'a, T:'a + Eq + Clone> {
  pub literal: T,
}

impl<'a, T: 'a + Eq + Clone> Parser<'a,  &'a [T], T> for LiteralParser<'a, T> {
  fn parse(&self, data: &'a[T]) -> ParseResult<'a, &'a[T], T> {
    if data.len() < 1 {
      return Err(format!("ran out of data"))
    }
    if data[0] == self.literal {
      Ok((data[0].clone(), data.slice_from(1)))
    } else {
      Err(format!("Literal mismatch"))
    }
  }
}

/// A string Parser that attempts to consume the given regex
pub struct RegexLiteralParser<'a> {
  pub regex: Regex,
}

impl<'a> Parser<'a, &'a str, ()> for RegexLiteralParser<'a> {
  fn parse(&self, data: &'a str) -> ParseResult<'a, &'a str, ()> {
    self.regex.find(data).map(|(_, e)| ((), data.slice_from(e))).ok_or(format!("regex literal match fail"))
  }
}

pub struct RegexCapturesParser<'a> {
  pub regex: Regex,
}
impl<'a> Parser<'a, &'a str, Captures<'a>> for RegexCapturesParser<'a> {
  fn parse(&self, data: &'a str) -> ParseResult<'a, &'a str, Captures<'a>> {
    match self.regex.captures(data) {
      Some(caps) => match caps.pos(0) {
        Some((_, e)) => Ok((caps, data.slice_from(e))),
        None => Err(format!("No Match"))
      },
      None => Err(format!("No Match"))
    }
  }
}

/// A slice Parser that matches against the first item in the slice
pub struct MatchParser<'a, I, O> {
  pub matcher: Box< Fn<(&'a I,), Result<O, String>> +'a>
}
impl<'a, I: Clone, O> Parser<'a, &'a [I], O> for MatchParser<'a, I, O> {
  fn parse(&self, data: &'a[I]) -> ParseResult<'a, &'a[I], O> {
    if data.len() < 1 {
      Err(format!("Unexpected End!"))
    } else {
      self.matcher.call((&data[0],)).map(|res| (res, data.slice_from(1)))
    }
  }
}

/// A Parser that will keep repeating the given parser until it returns an
/// error.  The accumulated results are returned.
pub struct RepParser<'a, I, O, P: Parser<'a, I, O>>{
  pub parser: P
}

impl<'a, I: Clone, O, P: Parser<'a, I, O>> Parser<'a, I, Vec<O>> for RepParser<'a, I, O, P> {
  fn parse(&self, data: I) -> ParseResult<'a, I, Vec<O>> {
    let mut remain = data;
    let mut v: Vec<O> = Vec::new();
    loop {
      match self.parser.parse(remain.clone()) {
        Ok((result, rest)) => {
          v.push(result);
          remain = rest;
        }
        Err(_) => {
          return Ok((v, remain));
        }
      }
    }
  }
}

/// A Parser that will repeatedly parse `rep` and `sep` in sequence until `sep`
/// returns an error.  The accumulated `rep` results are returned.  If `rep`
/// returns an error at any time, the error is escelated.
pub struct RepSepParser<'a, I, O, U, A: Parser<'a, I, O>, B: Parser<'a, I, U>> {
  pub rep: A,
  pub sep: B,
  pub min_reps: uint,
}
impl<'a, I: Clone, O, U, A: Parser<'a, I, O>, B: Parser<'a, I, U>> Parser<'a, I, Vec<O>> for RepSepParser<'a, I, O, U, A, B> {
  fn parse(&self, data: I) -> ParseResult<'a, I, Vec<O>> {
    let mut remain = data;
    let mut v: Vec<O> = Vec::new();
    loop {
      match self.rep.parse(remain) {
        Ok((result, rest)) => {
          v.push(result);
          match self.sep.parse(rest.clone()) {
            Ok((_, rest2)) => {
              remain = rest2
            }
            Err(_) => {
              if v.len() < self.min_reps {
                return Err(format!("Not enough reps: required {}, got {}", self.min_reps, v.len()))
              } else {
                return Ok((v, rest))
              }
            }
          }
        }
        Err(err) => {
          return Err(format!("Error on rep: {}", err));
        }
      }
    }
    unreachable!()
  }
}

/// A Parser that will combine two parsers into a tuple of results.  If either
/// parser returns an error, the error is escelated (ie partial successes are
/// not returned)
pub struct DualParser<'a, I, A, B, X: Parser<'a, I, A>, Y: Parser<'a, I, B>> {
  pub first: X,
  pub second: Y,
}

impl <'a, I, A, B, X: Parser<'a, I, A>, Y: Parser<'a, I, B>> Parser<'a, I, (A,B)> for DualParser<'a, I, A, B, X, Y> {
  fn parse(&self, data: I) -> ParseResult<'a, I, (A, B)> {
    match self.first.parse(data) {
      Ok((a, d2)) => match self.second.parse(d2) {
        Ok((b, remain)) => Ok(((a, b), remain)),
        Err(err) => Err(err)
      },
      Err(err) => Err(err)
    }
  }
}

//memoization coming sooooon
pub struct LazyParser<'a, I, O> {
  pub generator: Box<Fn<(), Box<Parser<'a, I, O> + 'a> + 'a> + 'a>
}
impl<'a, I, O> Parser<'a, I, O> for LazyParser<'a, I, O> {
  fn parse(&self, input: I) -> ParseResult<'a, I, O> {
    self.generator.call(()).parse(input)
  }
}

/// A parser that will attempt to parse using parser `a`, and then `b` if
/// the first fails.  The contained parsers are lazy so that we can support recursive
/// grammars.
///
/// In general, the "greedier" parser should be `a`.
pub struct OrParser<'a, I, O, A: Parser<'a, I, O>, B: Parser<'a, I, O>> {
  pub a: A,
  pub b: B,
}

impl<'a, I: Clone, O, A: Parser<'a, I, O>, B: Parser<'a, I, O>> Parser<'a, I, O> for OrParser<'a, I, O, A, B> {
  fn parse(&self, data: I) -> ParseResult<'a, I, O> {
    match self.a.parse(data.clone()) {
      Ok((a, d2)) => Ok((a, d2)),
      Err(_) => match self.b.parse(data.clone()) {
        Ok((b, remain)) => Ok((b, remain)),
        Err(err) => Err(err)
      }
    }
  }
}

/// A Parser that can map the successful result of a parser to another type
pub struct MapParser<'a, I, O, U, P: Parser<'a, I, O>> {
  pub parser: P,
  pub mapper: Box<Fn<(O,), U> + 'a>, //this has to be a &Fn and not a regular lambda since it must be immutable
}
impl<'a, I, O, U, P: Parser<'a, I, O>> Parser<'a, I, U> for MapParser<'a, I, O, U, P> {
  fn parse(&self, data: I) -> ParseResult<'a, I, U> {
    self.parser.parse(data).map(|(output, input)| ((self.mapper.call((output,)), input)))
  }
}

/// A Parser that will attempt to use a parser, returning an option of the contained parser's result
pub struct OptionParser<'a, I, O, P: Parser<'a, I, O>> {
  pub parser: P
}
impl<'a, I: Clone, O, P: Parser<'a, I, O>> Parser<'a, I, Option<O>> for OptionParser<'a, I, O, P> {
  fn parse(&self, data: I) -> ParseResult<'a, I, Option<O>> {
    match self.parser.parse(data.clone()) {
      Ok((result, rest))  => Ok((Some(result), rest)),
      Err(_)              => Ok((None, data)),
    }
  }
}

<commit_msg>Fix a syntax error<commit_after>
use std::ops::Fn;
use regex::{Captures, Regex};

/// A Parser is designed to take an input type and turn part of it into an
/// output type.  Probably in every real-life scenario, the input type is a
/// slice of some type.  Thus a parser returns an output type along with another
/// input type, which for slices would be the rest of the data.
pub trait Parser<'a, I, O> {

  fn parse(&self, data: I) -> ParseResult<'a, I, O>;

}

pub type ParseResult<'a, I:'a, O> = Result<(O, I), String>;

pub fn literal<'a, T:'a + Eq + Clone>(literal: T) -> LiteralParser<'a, T>{
  LiteralParser{literal: literal}
}

/// A parser that consumes one item in a slice and only returns ok if it equals
/// the given literal
pub struct LiteralParser<'a, T:'a + Eq + Clone> {
  pub literal: T,
}

impl<'a, T: 'a + Eq + Clone> Parser<'a,  &'a [T], T> for LiteralParser<'a, T> {
  fn parse(&self, data: &'a[T]) -> ParseResult<'a, &'a[T], T> {
    if data.len() < 1 {
      return Err(format!("ran out of data"))
    }
    if data[0] == self.literal {
      Ok((data[0].clone(), data.slice_from(1)))
    } else {
      Err(format!("Literal mismatch"))
    }
  }
}

/// A string Parser that attempts to consume the given regex
pub struct RegexLiteralParser<'a> {
  pub regex: Regex,
}

impl<'a> Parser<'a, &'a str, ()> for RegexLiteralParser<'a> {
  fn parse(&self, data: &'a str) -> ParseResult<'a, &'a str, ()> {
    self.regex.find(data).map(|(_, e)| ((), data.slice_from(e))).ok_or(format!("regex literal match fail"))
  }
}

pub struct RegexCapturesParser<'a> {
  pub regex: Regex,
}
impl<'a> Parser<'a, &'a str, Captures<'a>> for RegexCapturesParser<'a> {
  fn parse(&self, data: &'a str) -> ParseResult<'a, &'a str, Captures<'a>> {
    match self.regex.captures(data) {
      Some(caps) => match caps.pos(0) {
        Some((_, e)) => Ok((caps, data.slice_from(e))),
        None => Err(format!("No Match"))
      },
      None => Err(format!("No Match"))
    }
  }
}

/// A slice Parser that matches against the first item in the slice
pub struct MatchParser<'a, I, O> {
  pub matcher: Box< Fn<(&'a I,), Result<O, String>> +'a>
}
impl<'a, I: Clone, O> Parser<'a, &'a [I], O> for MatchParser<'a, I, O> {
  fn parse(&self, data: &'a[I]) -> ParseResult<'a, &'a[I], O> {
    if data.len() < 1 {
      Err(format!("Unexpected End!"))
    } else {
      self.matcher.call((&data[0],)).map(|res| (res, data.slice_from(1)))
    }
  }
}

/// A Parser that will keep repeating the given parser until it returns an
/// error.  The accumulated results are returned.
pub struct RepParser<'a, I, O, P: Parser<'a, I, O>>{
  pub parser: P
}

impl<'a, I: Clone, O, P: Parser<'a, I, O>> Parser<'a, I, Vec<O>> for RepParser<'a, I, O, P> {
  fn parse(&self, data: I) -> ParseResult<'a, I, Vec<O>> {
    let mut remain = data;
    let mut v: Vec<O> = Vec::new();
    loop {
      match self.parser.parse(remain.clone()) {
        Ok((result, rest)) => {
          v.push(result);
          remain = rest;
        }
        Err(_) => {
          return Ok((v, remain));
        }
      }
    }
  }
}

/// A Parser that will repeatedly parse `rep` and `sep` in sequence until `sep`
/// returns an error.  The accumulated `rep` results are returned.  If `rep`
/// returns an error at any time, the error is escelated.
pub struct RepSepParser<'a, I, O, U, A: Parser<'a, I, O>, B: Parser<'a, I, U>> {
  pub rep: A,
  pub sep: B,
  pub min_reps: uint,
}
impl<'a, I: Clone, O, U, A: Parser<'a, I, O>, B: Parser<'a, I, U>> Parser<'a, I, Vec<O>> for RepSepParser<'a, I, O, U, A, B> {
  fn parse(&self, data: I) -> ParseResult<'a, I, Vec<O>> {
    let mut remain = data;
    let mut v: Vec<O> = Vec::new();
    loop {
      match self.rep.parse(remain) {
        Ok((result, rest)) => {
          v.push(result);
          match self.sep.parse(rest.clone()) {
            Ok((_, rest2)) => {
              remain = rest2
            }
            Err(_) => {
              if v.len() < self.min_reps {
                return Err(format!("Not enough reps: required {}, got {}", self.min_reps, v.len()))
              } else {
                return Ok((v, rest))
              }
            }
          }
        }
        Err(err) => {
          return Err(format!("Error on rep: {}", err));
        }
      }
    }
    unreachable!()
  }
}

/// A Parser that will combine two parsers into a tuple of results.  If either
/// parser returns an error, the error is escelated (ie partial successes are
/// not returned)
pub struct DualParser<'a, I, A, B, X: Parser<'a, I, A>, Y: Parser<'a, I, B>> {
  pub first: X,
  pub second: Y,
}

impl <'a, I, A, B, X: Parser<'a, I, A>, Y: Parser<'a, I, B>> Parser<'a, I, (A,B)> for DualParser<'a, I, A, B, X, Y> {
  fn parse(&self, data: I) -> ParseResult<'a, I, (A, B)> {
    match self.first.parse(data) {
      Ok((a, d2)) => match self.second.parse(d2) {
        Ok((b, remain)) => Ok(((a, b), remain)),
        Err(err) => Err(err)
      },
      Err(err) => Err(err)
    }
  }
}

//memoization coming sooooon
pub struct LazyParser<'a, I, O> {
  pub generator: Box<Fn() -> Box<Parser<'a, I, O> + 'a> + 'a>
}
impl<'a, I, O> Parser<'a, I, O> for LazyParser<'a, I, O> {
  fn parse(&self, input: I) -> ParseResult<'a, I, O> {
    self.generator.call(()).parse(input)
  }
}

/// A parser that will attempt to parse using parser `a`, and then `b` if
/// the first fails.  The contained parsers are lazy so that we can support recursive
/// grammars.
///
/// In general, the "greedier" parser should be `a`.
pub struct OrParser<'a, I, O, A: Parser<'a, I, O>, B: Parser<'a, I, O>> {
  pub a: A,
  pub b: B,
}

impl<'a, I: Clone, O, A: Parser<'a, I, O>, B: Parser<'a, I, O>> Parser<'a, I, O> for OrParser<'a, I, O, A, B> {
  fn parse(&self, data: I) -> ParseResult<'a, I, O> {
    match self.a.parse(data.clone()) {
      Ok((a, d2)) => Ok((a, d2)),
      Err(_) => match self.b.parse(data.clone()) {
        Ok((b, remain)) => Ok((b, remain)),
        Err(err) => Err(err)
      }
    }
  }
}

/// A Parser that can map the successful result of a parser to another type
pub struct MapParser<'a, I, O, U, P: Parser<'a, I, O>> {
  pub parser: P,
  pub mapper: Box<Fn<(O,), U> + 'a>, //this has to be a &Fn and not a regular lambda since it must be immutable
}
impl<'a, I, O, U, P: Parser<'a, I, O>> Parser<'a, I, U> for MapParser<'a, I, O, U, P> {
  fn parse(&self, data: I) -> ParseResult<'a, I, U> {
    self.parser.parse(data).map(|(output, input)| ((self.mapper.call((output,)), input)))
  }
}

/// A Parser that will attempt to use a parser, returning an option of the contained parser's result
pub struct OptionParser<'a, I, O, P: Parser<'a, I, O>> {
  pub parser: P
}
impl<'a, I: Clone, O, P: Parser<'a, I, O>> Parser<'a, I, Option<O>> for OptionParser<'a, I, O, P> {
  fn parse(&self, data: I) -> ParseResult<'a, I, Option<O>> {
    match self.parser.parse(data.clone()) {
      Ok((result, rest))  => Ok((Some(result), rest)),
      Err(_)              => Ok((None, data)),
    }
  }
}

<|endoftext|>


<commit_before>use hir::HirDisplay;
use syntax::{ast, AstNode, SyntaxKind, SyntaxToken, TextRange, TextSize};

use crate::{AssistContext, AssistId, AssistKind, Assists};

// Assist: add_return_type
//
// Adds the return type to a function or closure inferred from its tail expression if it doesn't have a return
// type specified. This assists is useable in a functions or closures tail expression or return type position.
//
// ```
// fn foo() { 4$02i32 }
// ```
// ->
// ```
// fn foo() -> i32 { 42i32 }
// ```
pub(crate) fn add_return_type(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let (fn_type, tail_expr, builder_edit_pos, needs_whitespace) = extract_tail(ctx)?;
    let module = ctx.sema.scope(tail_expr.syntax()).module()?;
    let ty = ctx.sema.type_of_expr(&tail_expr)?.adjusted();
    if ty.is_unit() {
        return None;
    }
    let ty = ty.display_source_code(ctx.db(), module.into()).ok()?;

    acc.add(
        AssistId("add_return_type", AssistKind::RefactorRewrite),
        match fn_type {
            FnType::Function => "Add this function's return type",
            FnType::Closure { .. } => "Add this closure's return type",
        },
        tail_expr.syntax().text_range(),
        |builder| {
            let preceeding_whitespace = if needs_whitespace { " " } else { "" };

            match builder_edit_pos {
                InsertOrReplace::Insert(insert_pos) => {
                    builder.insert(insert_pos, &format!("{}-> {} ", preceeding_whitespace, ty))
                }
                InsertOrReplace::Replace(text_range) => {
                    builder.replace(text_range, &format!("{}-> {}", preceeding_whitespace, ty))
                }
            }
            if let FnType::Closure { wrap_expr: true } = fn_type {
                cov_mark::hit!(wrap_closure_non_block_expr);
                // `|x| x` becomes `|x| -> T x` which is invalid, so wrap it in a block
                builder.replace(tail_expr.syntax().text_range(), &format!("{{{}}}", tail_expr));
            }
        },
    )
}

enum InsertOrReplace {
    Insert(TextSize),
    Replace(TextRange),
}

/// Check the potentially already specified return type and reject it or turn it into a builder command
/// if allowed.
fn ret_ty_to_action(
    ret_ty: Option<ast::RetType>,
    insert_after: SyntaxToken,
) -> Option<(InsertOrReplace, bool)> {
    match ret_ty {
        Some(ret_ty) => match ret_ty.ty() {
            Some(ast::Type::InferType(_)) | None => {
                cov_mark::hit!(existing_infer_ret_type);
                cov_mark::hit!(existing_infer_ret_type_closure);
                Some((InsertOrReplace::Replace(ret_ty.syntax().text_range()), false))
            }
            _ => {
                cov_mark::hit!(existing_ret_type);
                cov_mark::hit!(existing_ret_type_closure);
                None
            }
        },
        None => {
            let insert_after_pos = insert_after.text_range().end();
            let (insert_pos, needs_whitespace) = match insert_after.next_token() {
                Some(it) if it.kind() == SyntaxKind::WHITESPACE => {
                    (insert_after_pos + TextSize::from(1), false)
                }
                _ => (insert_after_pos, true),
            };

            Some((InsertOrReplace::Insert(insert_pos), needs_whitespace))
        }
    }
}

enum FnType {
    Function,
    Closure { wrap_expr: bool },
}

fn extract_tail(ctx: &AssistContext) -> Option<(FnType, ast::Expr, InsertOrReplace, bool)> {
    let (fn_type, tail_expr, return_type_range, action, needs_whitespace) =
        if let Some(closure) = ctx.find_node_at_offset::<ast::ClosureExpr>() {
            let rpipe = closure.param_list()?.syntax().last_token()?;
            let rpipe_pos = rpipe.text_range().end();

            let (action, needs_whitespace) = ret_ty_to_action(closure.ret_type(), rpipe)?;

            let body = closure.body()?;
            let body_start = body.syntax().first_token()?.text_range().start();
            let (tail_expr, wrap_expr) = match body {
                ast::Expr::BlockExpr(block) => (block.tail_expr()?, false),
                body => (body, true),
            };

            let ret_range = TextRange::new(rpipe_pos, body_start);
            (FnType::Closure { wrap_expr }, tail_expr, ret_range, action, needs_whitespace)
        } else {
            let func = ctx.find_node_at_offset::<ast::Fn>()?;

            let rparen = func.param_list()?.r_paren_token()?;
            let rparen_pos = rparen.text_range().end();
            let (action, needs_whitespace) = ret_ty_to_action(func.ret_type(), rparen)?;

            let body = func.body()?;
            let stmt_list = body.stmt_list()?;
            let tail_expr = stmt_list.tail_expr()?;

            let ret_range_end = stmt_list.l_curly_token()?.text_range().start();
            let ret_range = TextRange::new(rparen_pos, ret_range_end);
            (FnType::Function, tail_expr, ret_range, action, needs_whitespace)
        };
    let range = ctx.selection_trimmed();
    if return_type_range.contains_range(range) {
        cov_mark::hit!(cursor_in_ret_position);
        cov_mark::hit!(cursor_in_ret_position_closure);
    } else if tail_expr.syntax().text_range().contains_range(range) {
        cov_mark::hit!(cursor_on_tail);
        cov_mark::hit!(cursor_on_tail_closure);
    } else {
        return None;
    }
    Some((fn_type, tail_expr, action, needs_whitespace))
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn infer_return_type_specified_inferred() {
        cov_mark::check!(existing_infer_ret_type);
        check_assist(
            add_return_type,
            r#"fn foo() -> $0_ {
    45
}"#,
            r#"fn foo() -> i32 {
    45
}"#,
        );
    }

    #[test]
    fn infer_return_type_specified_inferred_closure() {
        cov_mark::check!(existing_infer_ret_type_closure);
        check_assist(
            add_return_type,
            r#"fn foo() {
    || -> _ {$045};
}"#,
            r#"fn foo() {
    || -> i32 {45};
}"#,
        );
    }

    #[test]
    fn infer_return_type_cursor_at_return_type_pos() {
        cov_mark::check!(cursor_in_ret_position);
        check_assist(
            add_return_type,
            r#"fn foo() $0{
    45
}"#,
            r#"fn foo() -> i32 {
    45
}"#,
        );
    }

    #[test]
    fn infer_return_type_cursor_at_return_type_pos_closure() {
        cov_mark::check!(cursor_in_ret_position_closure);
        check_assist(
            add_return_type,
            r#"fn foo() {
    || $045
}"#,
            r#"fn foo() {
    || -> i32 {45}
}"#,
        );
    }

    #[test]
    fn infer_return_type() {
        cov_mark::check!(cursor_on_tail);
        check_assist(
            add_return_type,
            r#"fn foo() {
    45$0
}"#,
            r#"fn foo() -> i32 {
    45
}"#,
        );
    }

    #[test]
    fn infer_return_type_no_whitespace() {
        check_assist(
            add_return_type,
            r#"fn foo(){
    45$0
}"#,
            r#"fn foo() -> i32 {
    45
}"#,
        );
    }

    #[test]
    fn infer_return_type_nested() {
        check_assist(
            add_return_type,
            r#"fn foo() {
    if true {
        3$0
    } else {
        5
    }
}"#,
            r#"fn foo() -> i32 {
    if true {
        3
    } else {
        5
    }
}"#,
        );
    }

    #[test]
    fn not_applicable_ret_type_specified() {
        cov_mark::check!(existing_ret_type);
        check_assist_not_applicable(
            add_return_type,
            r#"fn foo() -> i32 {
    ( 45$0 + 32 ) * 123
}"#,
        );
    }

    #[test]
    fn not_applicable_non_tail_expr() {
        check_assist_not_applicable(
            add_return_type,
            r#"fn foo() {
    let x = $03;
    ( 45 + 32 ) * 123
}"#,
        );
    }

    #[test]
    fn not_applicable_unit_return_type() {
        check_assist_not_applicable(
            add_return_type,
            r#"fn foo() {
    ($0)
}"#,
        );
    }

    #[test]
    fn infer_return_type_closure_block() {
        cov_mark::check!(cursor_on_tail_closure);
        check_assist(
            add_return_type,
            r#"fn foo() {
    |x: i32| {
        x$0
    };
}"#,
            r#"fn foo() {
    |x: i32| -> i32 {
        x
    };
}"#,
        );
    }

    #[test]
    fn infer_return_type_closure() {
        check_assist(
            add_return_type,
            r#"fn foo() {
    |x: i32| { x$0 };
}"#,
            r#"fn foo() {
    |x: i32| -> i32 { x };
}"#,
        );
    }

    #[test]
    fn infer_return_type_closure_no_whitespace() {
        check_assist(
            add_return_type,
            r#"fn foo() {
    |x: i32|{ x$0 };
}"#,
            r#"fn foo() {
    |x: i32| -> i32 { x };
}"#,
        );
    }

    #[test]
    fn infer_return_type_closure_wrap() {
        cov_mark::check!(wrap_closure_non_block_expr);
        check_assist(
            add_return_type,
            r#"fn foo() {
    |x: i32| x$0;
}"#,
            r#"fn foo() {
    |x: i32| -> i32 {x};
}"#,
        );
    }

    #[test]
    fn infer_return_type_nested_closure() {
        check_assist(
            add_return_type,
            r#"fn foo() {
    || {
        if true {
            3$0
        } else {
            5
        }
    }
}"#,
            r#"fn foo() {
    || -> i32 {
        if true {
            3
        } else {
            5
        }
    }
}"#,
        );
    }

    #[test]
    fn not_applicable_ret_type_specified_closure() {
        cov_mark::check!(existing_ret_type_closure);
        check_assist_not_applicable(
            add_return_type,
            r#"fn foo() {
    || -> i32 { 3$0 }
}"#,
        );
    }

    #[test]
    fn not_applicable_non_tail_expr_closure() {
        check_assist_not_applicable(
            add_return_type,
            r#"fn foo() {
    || -> i32 {
        let x = 3$0;
        6
    }
}"#,
        );
    }
}
<commit_msg>refactor(assist/add_return_type): avoid threading `needs_whitespace`<commit_after>use hir::HirDisplay;
use syntax::{ast, AstNode, SyntaxKind, SyntaxToken, TextRange, TextSize};

use crate::{AssistContext, AssistId, AssistKind, Assists};

// Assist: add_return_type
//
// Adds the return type to a function or closure inferred from its tail expression if it doesn't have a return
// type specified. This assists is useable in a functions or closures tail expression or return type position.
//
// ```
// fn foo() { 4$02i32 }
// ```
// ->
// ```
// fn foo() -> i32 { 42i32 }
// ```
pub(crate) fn add_return_type(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let (fn_type, tail_expr, builder_edit_pos) = extract_tail(ctx)?;
    let module = ctx.sema.scope(tail_expr.syntax()).module()?;
    let ty = ctx.sema.type_of_expr(&tail_expr)?.adjusted();
    if ty.is_unit() {
        return None;
    }
    let ty = ty.display_source_code(ctx.db(), module.into()).ok()?;

    acc.add(
        AssistId("add_return_type", AssistKind::RefactorRewrite),
        match fn_type {
            FnType::Function => "Add this function's return type",
            FnType::Closure { .. } => "Add this closure's return type",
        },
        tail_expr.syntax().text_range(),
        |builder| {
            match builder_edit_pos {
                InsertOrReplace::Insert(insert_pos, needs_whitespace) => {
                    let preceeding_whitespace = if needs_whitespace { " " } else { "" };
                    builder.insert(insert_pos, &format!("{}-> {} ", preceeding_whitespace, ty))
                }
                InsertOrReplace::Replace(text_range) => {
                    builder.replace(text_range, &format!("-> {}", ty))
                }
            }
            if let FnType::Closure { wrap_expr: true } = fn_type {
                cov_mark::hit!(wrap_closure_non_block_expr);
                // `|x| x` becomes `|x| -> T x` which is invalid, so wrap it in a block
                builder.replace(tail_expr.syntax().text_range(), &format!("{{{}}}", tail_expr));
            }
        },
    )
}

enum InsertOrReplace {
    Insert(TextSize, bool),
    Replace(TextRange),
}

/// Check the potentially already specified return type and reject it or turn it into a builder command
/// if allowed.
fn ret_ty_to_action(
    ret_ty: Option<ast::RetType>,
    insert_after: SyntaxToken,
) -> Option<InsertOrReplace> {
    match ret_ty {
        Some(ret_ty) => match ret_ty.ty() {
            Some(ast::Type::InferType(_)) | None => {
                cov_mark::hit!(existing_infer_ret_type);
                cov_mark::hit!(existing_infer_ret_type_closure);
                Some(InsertOrReplace::Replace(ret_ty.syntax().text_range()))
            }
            _ => {
                cov_mark::hit!(existing_ret_type);
                cov_mark::hit!(existing_ret_type_closure);
                None
            }
        },
        None => {
            let insert_after_pos = insert_after.text_range().end();
            let (insert_pos, needs_whitespace) = match insert_after.next_token() {
                Some(it) if it.kind() == SyntaxKind::WHITESPACE => {
                    (insert_after_pos + TextSize::from(1), false)
                }
                _ => (insert_after_pos, true),
            };

            Some(InsertOrReplace::Insert(insert_pos, needs_whitespace))
        }
    }
}

enum FnType {
    Function,
    Closure { wrap_expr: bool },
}

fn extract_tail(ctx: &AssistContext) -> Option<(FnType, ast::Expr, InsertOrReplace)> {
    let (fn_type, tail_expr, return_type_range, action) =
        if let Some(closure) = ctx.find_node_at_offset::<ast::ClosureExpr>() {
            let rpipe = closure.param_list()?.syntax().last_token()?;
            let rpipe_pos = rpipe.text_range().end();

            let action = ret_ty_to_action(closure.ret_type(), rpipe)?;

            let body = closure.body()?;
            let body_start = body.syntax().first_token()?.text_range().start();
            let (tail_expr, wrap_expr) = match body {
                ast::Expr::BlockExpr(block) => (block.tail_expr()?, false),
                body => (body, true),
            };

            let ret_range = TextRange::new(rpipe_pos, body_start);
            (FnType::Closure { wrap_expr }, tail_expr, ret_range, action)
        } else {
            let func = ctx.find_node_at_offset::<ast::Fn>()?;

            let rparen = func.param_list()?.r_paren_token()?;
            let rparen_pos = rparen.text_range().end();
            let action = ret_ty_to_action(func.ret_type(), rparen)?;

            let body = func.body()?;
            let stmt_list = body.stmt_list()?;
            let tail_expr = stmt_list.tail_expr()?;

            let ret_range_end = stmt_list.l_curly_token()?.text_range().start();
            let ret_range = TextRange::new(rparen_pos, ret_range_end);
            (FnType::Function, tail_expr, ret_range, action)
        };
    let range = ctx.selection_trimmed();
    if return_type_range.contains_range(range) {
        cov_mark::hit!(cursor_in_ret_position);
        cov_mark::hit!(cursor_in_ret_position_closure);
    } else if tail_expr.syntax().text_range().contains_range(range) {
        cov_mark::hit!(cursor_on_tail);
        cov_mark::hit!(cursor_on_tail_closure);
    } else {
        return None;
    }
    Some((fn_type, tail_expr, action))
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn infer_return_type_specified_inferred() {
        cov_mark::check!(existing_infer_ret_type);
        check_assist(
            add_return_type,
            r#"fn foo() -> $0_ {
    45
}"#,
            r#"fn foo() -> i32 {
    45
}"#,
        );
    }

    #[test]
    fn infer_return_type_specified_inferred_closure() {
        cov_mark::check!(existing_infer_ret_type_closure);
        check_assist(
            add_return_type,
            r#"fn foo() {
    || -> _ {$045};
}"#,
            r#"fn foo() {
    || -> i32 {45};
}"#,
        );
    }

    #[test]
    fn infer_return_type_cursor_at_return_type_pos() {
        cov_mark::check!(cursor_in_ret_position);
        check_assist(
            add_return_type,
            r#"fn foo() $0{
    45
}"#,
            r#"fn foo() -> i32 {
    45
}"#,
        );
    }

    #[test]
    fn infer_return_type_cursor_at_return_type_pos_closure() {
        cov_mark::check!(cursor_in_ret_position_closure);
        check_assist(
            add_return_type,
            r#"fn foo() {
    || $045
}"#,
            r#"fn foo() {
    || -> i32 {45}
}"#,
        );
    }

    #[test]
    fn infer_return_type() {
        cov_mark::check!(cursor_on_tail);
        check_assist(
            add_return_type,
            r#"fn foo() {
    45$0
}"#,
            r#"fn foo() -> i32 {
    45
}"#,
        );
    }

    #[test]
    fn infer_return_type_no_whitespace() {
        check_assist(
            add_return_type,
            r#"fn foo(){
    45$0
}"#,
            r#"fn foo() -> i32 {
    45
}"#,
        );
    }

    #[test]
    fn infer_return_type_nested() {
        check_assist(
            add_return_type,
            r#"fn foo() {
    if true {
        3$0
    } else {
        5
    }
}"#,
            r#"fn foo() -> i32 {
    if true {
        3
    } else {
        5
    }
}"#,
        );
    }

    #[test]
    fn not_applicable_ret_type_specified() {
        cov_mark::check!(existing_ret_type);
        check_assist_not_applicable(
            add_return_type,
            r#"fn foo() -> i32 {
    ( 45$0 + 32 ) * 123
}"#,
        );
    }

    #[test]
    fn not_applicable_non_tail_expr() {
        check_assist_not_applicable(
            add_return_type,
            r#"fn foo() {
    let x = $03;
    ( 45 + 32 ) * 123
}"#,
        );
    }

    #[test]
    fn not_applicable_unit_return_type() {
        check_assist_not_applicable(
            add_return_type,
            r#"fn foo() {
    ($0)
}"#,
        );
    }

    #[test]
    fn infer_return_type_closure_block() {
        cov_mark::check!(cursor_on_tail_closure);
        check_assist(
            add_return_type,
            r#"fn foo() {
    |x: i32| {
        x$0
    };
}"#,
            r#"fn foo() {
    |x: i32| -> i32 {
        x
    };
}"#,
        );
    }

    #[test]
    fn infer_return_type_closure() {
        check_assist(
            add_return_type,
            r#"fn foo() {
    |x: i32| { x$0 };
}"#,
            r#"fn foo() {
    |x: i32| -> i32 { x };
}"#,
        );
    }

    #[test]
    fn infer_return_type_closure_no_whitespace() {
        check_assist(
            add_return_type,
            r#"fn foo() {
    |x: i32|{ x$0 };
}"#,
            r#"fn foo() {
    |x: i32| -> i32 { x };
}"#,
        );
    }

    #[test]
    fn infer_return_type_closure_wrap() {
        cov_mark::check!(wrap_closure_non_block_expr);
        check_assist(
            add_return_type,
            r#"fn foo() {
    |x: i32| x$0;
}"#,
            r#"fn foo() {
    |x: i32| -> i32 {x};
}"#,
        );
    }

    #[test]
    fn infer_return_type_nested_closure() {
        check_assist(
            add_return_type,
            r#"fn foo() {
    || {
        if true {
            3$0
        } else {
            5
        }
    }
}"#,
            r#"fn foo() {
    || -> i32 {
        if true {
            3
        } else {
            5
        }
    }
}"#,
        );
    }

    #[test]
    fn not_applicable_ret_type_specified_closure() {
        cov_mark::check!(existing_ret_type_closure);
        check_assist_not_applicable(
            add_return_type,
            r#"fn foo() {
    || -> i32 { 3$0 }
}"#,
        );
    }

    #[test]
    fn not_applicable_non_tail_expr_closure() {
        check_assist_not_applicable(
            add_return_type,
            r#"fn foo() {
    || -> i32 {
        let x = 3$0;
        6
    }
}"#,
        );
    }
}
<|endoftext|>


<commit_before>use crate::{api::prelude::*, proc_macros::*};

enum Action {
    ToggleSelection,
}

/// The `SelectionBehaviorState` handles the `SelectionBehavior` widget.
#[derive(Default, AsAny)]
pub struct SelectionBehaviorState {
    target: Entity,
}

impl State for SelectionBehaviorState {
    fn init(&mut self, _: &mut Registry, ctx: &mut Context) {
        self.target = (*ctx.widget().get::<u32>("target")).into();
        toggle_flag("selected", &mut ctx.get_widget(self.target));
        ctx.get_widget(self.target).update(false);
    }

    fn messages(
        &mut self,
        mut messages: MessageReader,
        _registry: &mut Registry,
        ctx: &mut Context,
    ) {
        for message in messages.read::<Action>() {
            match message {
                Action::ToggleSelection => {
                    let selected = *ctx.get_widget(self.target).get::<bool>("selected");
                    ctx.get_widget(self.target).set("selected", !selected);
                    toggle_flag("selected", &mut ctx.get_widget(self.target));
                    ctx.get_widget(self.target).update(false);
                }
            };
        }
    }
}

widget!(
    /// The `SelectionBehavior` widget is used to handle internal the pressed behavior of a widget.
    ///
    /// **style:** `check-box`
    SelectionBehavior<SelectionBehaviorState>: MouseHandler {
        /// Sets or shares the target of the behavior.
        target: u32,

        /// Sets or shares the selected property.
        selected: bool,

        /// Sets the parent id.
        parent: u32
    }
);

impl Template for SelectionBehavior {
    fn template(self, id: Entity, _: &mut BuildContext) -> Self {
        self.name("SelectionBehavior")
            .selected(true)
            .on_click(move |ctx, _| {
                ctx.send_message(Action::ToggleSelection, id);
                false
            })
    }
}
<commit_msg>more messages<commit_after>use crate::{api::prelude::*, proc_macros::*};

/// Used for selection.
pub enum SelectionAction {
    ToggleSelection,
}

/// The `SelectionBehaviorState` handles the `SelectionBehavior` widget.
#[derive(Default, AsAny)]
pub struct SelectionBehaviorState {
    target: Entity,
}

impl State for SelectionBehaviorState {
    fn init(&mut self, _: &mut Registry, ctx: &mut Context) {
        self.target = (*ctx.widget().get::<u32>("target")).into();
        toggle_flag("selected", &mut ctx.get_widget(self.target));
        ctx.get_widget(self.target).update(false);
    }

    fn messages(
        &mut self,
        mut messages: MessageReader,
        _registry: &mut Registry,
        ctx: &mut Context,
    ) {
        for message in messages.read::<SelectionAction>() {
            match message {
                SelectionAction::ToggleSelection => {
                    let selected = *ctx.get_widget(self.target).get::<bool>("selected");
                    ctx.get_widget(self.target).set("selected", !selected);
                    toggle_flag("selected", &mut ctx.get_widget(self.target));
                    ctx.get_widget(self.target).update(false);
                }
            };
        }
    }
}

widget!(
    /// The `SelectionBehavior` widget is used to handle internal the pressed behavior of a widget.
    ///
    /// **style:** `check-box`
    SelectionBehavior<SelectionBehaviorState>: MouseHandler {
        /// Sets or shares the target of the behavior.
        target: u32,

        /// Sets or shares the selected property.
        selected: bool,

        /// Sets the parent id.
        parent: u32
    }
);

impl Template for SelectionBehavior {
    fn template(self, id: Entity, _: &mut BuildContext) -> Self {
        self.name("SelectionBehavior")
            .selected(true)
            .on_click(move |ctx, _| {
                ctx.send_message(SelectionAction::ToggleSelection, id);
                false
            })
    }
}
<|endoftext|>


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

#[derive(Debug, Clone)]
pub struct Link(String);

impl Deref for Link {
    type Target = String;

    fn deref(&self) => &String {
        &self.0
    }

}

impl DerefMut for Link {

    fn deref_mut(&mut self) => &mut String {
        &mut self.0
    }

}
<commit_msg>fixup! Add Link type<commit_after>use std::ops::{Deref, DerefMut};

#[derive(Debug, Clone)]
pub struct Link(String);

impl Deref for Link {
    type Target = String;

    fn deref(&self) -> &String {
        &self.0
    }

}

impl DerefMut for Link {

    fn deref_mut(&mut self) -> &mut String {
        &mut self.0
    }

}
<|endoftext|>