<commit_before>use alloc::boxed::Box;
use alloc::collections::BTreeMap;
use alloc::vec::Vec;
use core::{cmp, str};
use core::sync::atomic::{AtomicUsize, Ordering};
use spin::RwLock;

use crate::syscall::data::Stat;
use crate::syscall::error::{Error, EBADF, EINVAL, ENOENT, Result};
use crate::syscall::flag::{MODE_DIR, MODE_FILE, SEEK_CUR, SEEK_END, SEEK_SET};
use crate::syscall::scheme::Scheme;

mod context;
mod cpu;
mod exe;
mod iostat;
mod log;
mod scheme;
mod scheme_num;
mod syscall;
mod uname;

struct Handle {
    path: &'static [u8],
    data: Vec<u8>,
    mode: u16,
    seek: usize
}

type SysFn = dyn Fn() -> Result<Vec<u8>> + Send + Sync;

/// System information scheme
pub struct SysScheme {
    next_id: AtomicUsize,
    files: BTreeMap<&'static [u8], Box<SysFn>>,
    handles: RwLock<BTreeMap<usize, Handle>>
}

impl SysScheme {
    pub fn new() -> SysScheme {
        let mut files: BTreeMap<&'static [u8], Box<SysFn>> = BTreeMap::new();

        files.insert(b"context", Box::new(context::resource));
        files.insert(b"cpu", Box::new(cpu::resource));
        files.insert(b"exe", Box::new(exe::resource));
        files.insert(b"iostat", Box::new(iostat::resource));
        files.insert(b"log", Box::new(log::resource));
        files.insert(b"scheme", Box::new(scheme::resource));
        files.insert(b"scheme_num", Box::new(scheme_num::resource));
        files.insert(b"syscall", Box::new(syscall::resource));
        files.insert(b"uname", Box::new(uname::resource));

        SysScheme {
            next_id: AtomicUsize::new(0),
            files,
            handles: RwLock::new(BTreeMap::new())
        }
    }
}

impl Scheme for SysScheme {
    fn open(&self, path: &[u8], _flags: usize, _uid: u32, _gid: u32) -> Result<usize> {
        let path_utf8 = str::from_utf8(path).or(Err(Error::new(ENOENT)))?;
        let path_trimmed = path_utf8.trim_matches('/');

        if path_trimmed.is_empty() {
            let mut data = Vec::new();
            for entry in self.files.iter() {
                if ! data.is_empty() {
                    data.push(b'\n');
                }
                data.extend_from_slice(entry.0);
            }

            let id = self.next_id.fetch_add(1, Ordering::SeqCst);
            self.handles.write().insert(id, Handle {
                path: b"",
                data,
                mode: MODE_DIR | 0o444,
                seek: 0
            });
            return Ok(id)
        } else {
            //Have to iterate to get the path without allocation
            for entry in self.files.iter() {
                if entry.0 == &path_trimmed.as_bytes() {
                    let id = self.next_id.fetch_add(1, Ordering::SeqCst);
                    self.handles.write().insert(id, Handle {
                        path: entry.0,
                        data: entry.1()?,
                        mode: MODE_FILE | 0o444,
                        seek: 0
                    });
                    return Ok(id)
                }
            }
        }

        Err(Error::new(ENOENT))
    }

    fn read(&self, id: usize, buffer: &mut [u8]) -> Result<usize> {
        let mut handles = self.handles.write();
        let handle = handles.get_mut(&id).ok_or(Error::new(EBADF))?;

        let mut i = 0;
        while i < buffer.len() && handle.seek < handle.data.len() {
            buffer[i] = handle.data[handle.seek];
            i += 1;
            handle.seek += 1;
        }

        Ok(i)
    }

    fn seek(&self, id: usize, pos: usize, whence: usize) -> Result<usize> {
        let mut handles = self.handles.write();
        let handle = handles.get_mut(&id).ok_or(Error::new(EBADF))?;

        handle.seek = match whence {
            SEEK_SET => cmp::min(handle.data.len(), pos),
            SEEK_CUR => cmp::max(0, cmp::min(handle.data.len() as isize, handle.seek as isize + pos as isize)) as usize,
            SEEK_END => cmp::max(0, cmp::min(handle.data.len() as isize, handle.data.len() as isize + pos as isize)) as usize,
            _ => return Err(Error::new(EINVAL))
        };

        Ok(handle.seek)
    }

    fn fpath(&self, id: usize, buf: &mut [u8]) -> Result<usize> {
        let handles = self.handles.read();
        let handle = handles.get(&id).ok_or(Error::new(EBADF))?;

        let mut i = 0;
        let scheme_path = b"sys:";
        while i < buf.len() && i < scheme_path.len() {
            buf[i] = scheme_path[i];
            i += 1;
        }

        let mut j = 0;
        while i < buf.len() && j < handle.path.len() {
            buf[i] = handle.path[j];
            i += 1;
            j += 1;
        }

        Ok(i)
    }

    fn fstat(&self, id: usize, stat: &mut Stat) -> Result<usize> {
        let handles = self.handles.read();
        let handle = handles.get(&id).ok_or(Error::new(EBADF))?;

        stat.st_mode = handle.mode;
        stat.st_uid = 0;
        stat.st_gid = 0;
        stat.st_size = handle.data.len() as u64;

        Ok(0)
    }

    fn fsync(&self, _id: usize) -> Result<usize> {
        Ok(0)
    }

    fn close(&self, id: usize) -> Result<usize> {
        self.handles.write().remove(&id).ok_or(Error::new(EBADF)).and(Ok(0))
    }
}
<commit_msg>Fix deadlock in sys: scheme<commit_after>use alloc::boxed::Box;
use alloc::collections::BTreeMap;
use alloc::vec::Vec;
use core::{cmp, str};
use core::sync::atomic::{AtomicUsize, Ordering};
use spin::RwLock;

use crate::syscall::data::Stat;
use crate::syscall::error::{Error, EBADF, EINVAL, ENOENT, Result};
use crate::syscall::flag::{MODE_DIR, MODE_FILE, SEEK_CUR, SEEK_END, SEEK_SET};
use crate::syscall::scheme::Scheme;

mod context;
mod cpu;
mod exe;
mod iostat;
mod log;
mod scheme;
mod scheme_num;
mod syscall;
mod uname;

struct Handle {
    path: &'static [u8],
    data: Vec<u8>,
    mode: u16,
    seek: usize
}

type SysFn = dyn Fn() -> Result<Vec<u8>> + Send + Sync;

/// System information scheme
pub struct SysScheme {
    next_id: AtomicUsize,
    files: BTreeMap<&'static [u8], Box<SysFn>>,
    handles: RwLock<BTreeMap<usize, Handle>>
}

impl SysScheme {
    pub fn new() -> SysScheme {
        let mut files: BTreeMap<&'static [u8], Box<SysFn>> = BTreeMap::new();

        files.insert(b"context", Box::new(context::resource));
        files.insert(b"cpu", Box::new(cpu::resource));
        files.insert(b"exe", Box::new(exe::resource));
        files.insert(b"iostat", Box::new(iostat::resource));
        files.insert(b"log", Box::new(log::resource));
        files.insert(b"scheme", Box::new(scheme::resource));
        files.insert(b"scheme_num", Box::new(scheme_num::resource));
        files.insert(b"syscall", Box::new(syscall::resource));
        files.insert(b"uname", Box::new(uname::resource));

        SysScheme {
            next_id: AtomicUsize::new(0),
            files,
            handles: RwLock::new(BTreeMap::new())
        }
    }
}

impl Scheme for SysScheme {
    fn open(&self, path: &[u8], _flags: usize, _uid: u32, _gid: u32) -> Result<usize> {
        let path_utf8 = str::from_utf8(path).or(Err(Error::new(ENOENT)))?;
        let path_trimmed = path_utf8.trim_matches('/');

        if path_trimmed.is_empty() {
            let mut data = Vec::new();
            for entry in self.files.iter() {
                if ! data.is_empty() {
                    data.push(b'\n');
                }
                data.extend_from_slice(entry.0);
            }

            let id = self.next_id.fetch_add(1, Ordering::SeqCst);
            self.handles.write().insert(id, Handle {
                path: b"",
                data,
                mode: MODE_DIR | 0o444,
                seek: 0
            });
            return Ok(id)
        } else {
            //Have to iterate to get the path without allocation
            for entry in self.files.iter() {
                if entry.0 == &path_trimmed.as_bytes() {
                    let id = self.next_id.fetch_add(1, Ordering::SeqCst);
                    let data = entry.1()?;
                    self.handles.write().insert(id, Handle {
                        path: entry.0,
                        data: data,
                        mode: MODE_FILE | 0o444,
                        seek: 0
                    });
                    return Ok(id)
                }
            }
        }

        Err(Error::new(ENOENT))
    }

    fn read(&self, id: usize, buffer: &mut [u8]) -> Result<usize> {
        let mut handles = self.handles.write();
        let handle = handles.get_mut(&id).ok_or(Error::new(EBADF))?;

        let mut i = 0;
        while i < buffer.len() && handle.seek < handle.data.len() {
            buffer[i] = handle.data[handle.seek];
            i += 1;
            handle.seek += 1;
        }

        Ok(i)
    }

    fn seek(&self, id: usize, pos: usize, whence: usize) -> Result<usize> {
        let mut handles = self.handles.write();
        let handle = handles.get_mut(&id).ok_or(Error::new(EBADF))?;

        handle.seek = match whence {
            SEEK_SET => cmp::min(handle.data.len(), pos),
            SEEK_CUR => cmp::max(0, cmp::min(handle.data.len() as isize, handle.seek as isize + pos as isize)) as usize,
            SEEK_END => cmp::max(0, cmp::min(handle.data.len() as isize, handle.data.len() as isize + pos as isize)) as usize,
            _ => return Err(Error::new(EINVAL))
        };

        Ok(handle.seek)
    }

    fn fpath(&self, id: usize, buf: &mut [u8]) -> Result<usize> {
        let handles = self.handles.read();
        let handle = handles.get(&id).ok_or(Error::new(EBADF))?;

        let mut i = 0;
        let scheme_path = b"sys:";
        while i < buf.len() && i < scheme_path.len() {
            buf[i] = scheme_path[i];
            i += 1;
        }

        let mut j = 0;
        while i < buf.len() && j < handle.path.len() {
            buf[i] = handle.path[j];
            i += 1;
            j += 1;
        }

        Ok(i)
    }

    fn fstat(&self, id: usize, stat: &mut Stat) -> Result<usize> {
        let handles = self.handles.read();
        let handle = handles.get(&id).ok_or(Error::new(EBADF))?;

        stat.st_mode = handle.mode;
        stat.st_uid = 0;
        stat.st_gid = 0;
        stat.st_size = handle.data.len() as u64;

        Ok(0)
    }

    fn fsync(&self, _id: usize) -> Result<usize> {
        Ok(0)
    }

    fn close(&self, id: usize) -> Result<usize> {
        self.handles.write().remove(&id).ok_or(Error::new(EBADF)).and(Ok(0))
    }
}
<|endoftext|>


<commit_before>extern crate ndarray;
extern crate numpy;
extern crate pyo3;

use ndarray::*;
use numpy::*;
use pyo3::prelude::*;

#[test]
fn new() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let n = 3;
    let m = 5;
    let arr = PyArray::<f64>::new(gil.python(), &np, &[n, m]);
    assert!(arr.ndim() == 2);
    assert!(arr.dims() == [n, m]);
    assert!(arr.strides() == [m as isize * 8, 8]);
}

#[test]
fn zeros() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let n = 3;
    let m = 5;
    let arr = PyArray::<f64>::zeros(gil.python(), &np, &[n, m], NPY_CORDER);
    assert!(arr.ndim() == 2);
    assert!(arr.dims() == [n, m]);
    assert!(arr.strides() == [m as isize * 8, 8]);

    let arr = PyArray::<f64>::zeros(gil.python(), &np, &[n, m], NPY_FORTRANORDER);
    assert!(arr.ndim() == 2);
    assert!(arr.dims() == [n, m]);
    assert!(arr.strides() == [8, n as isize * 8]);
}

#[test]
fn arange() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let arr = PyArray::<f64>::arange(gil.python(), &np, 0.0, 1.0, 0.1);
    println!("ndim = {:?}", arr.ndim());
    println!("dims = {:?}", arr.dims());
    println!("array = {:?}", arr.as_slice().unwrap());
}

#[test]
fn as_array() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let arr = PyArray::<f64>::zeros(gil.python(), &np, &[3, 2, 4], NPY_CORDER);
    let a = arr.as_array().unwrap();
    assert_eq!(arr.shape(), a.shape());
    assert_eq!(
        arr.strides().iter().map(|x| x / 8).collect::<Vec<_>>(),
        a.strides()
    );
}

#[test]
fn into_pyarray_vec() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();

    let a = vec![1, 2, 3];
    let arr = a.into_pyarray(gil.python(), &np);
    println!("arr.shape = {:?}", arr.shape());
    println!("arr = {:?}", arr.as_slice().unwrap());
    assert_eq!(arr.shape(), [3]);
}

#[test]
fn into_pyarray_array() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();

    let a = Array3::<f64>::zeros((3, 4, 2));
    let shape = a.shape().iter().cloned().collect::<Vec<_>>();
    let strides = a.strides().iter().map(|d| d * 8).collect::<Vec<_>>();
    println!("a.shape   = {:?}", a.shape());
    println!("a.strides = {:?}", a.strides());
    let pa = a.into_pyarray(gil.python(), &np);
    println!("pa.shape   = {:?}", pa.shape());
    println!("pa.strides = {:?}", pa.strides());
    assert_eq!(pa.shape(), shape.as_slice());
    assert_eq!(pa.strides(), strides.as_slice());
}

#[test]
fn iter_to_pyarray() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let arr = PyArray::from_iter(gil.python(), &np, (0..10).map(|x| x * x));
    assert_eq!(arr.as_slice().unwrap(), &[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]);
}

#[test]
fn is_instance() {
    let gil = pyo3::Python::acquire_gil();
    let py = gil.python();
    let np = PyArrayModule::import(py).unwrap();
    let arr = PyArray::<f64>::new(gil.python(), &np, &[3, 5]);
    assert!(py.is_instance::<PyArray<f64>, _>(&arr).unwrap());
    assert!(!py.is_instance::<pyo3::PyList, _>(&arr).unwrap());
}

#[test]
fn from_vec2() {
    let vec2 = vec![vec![1, 2, 3]; 2];
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let pyarray = PyArray::from_vec2(gil.python(), &np, &vec2).unwrap();
    assert_eq!(
        pyarray.as_array().unwrap(),
        array![[1, 2, 3], [1, 2, 3]].into_dyn()
    );
    assert!(PyArray::from_vec2(gil.python(), &np, &vec![vec![1], vec![2, 3]]).is_err());
}

#[test]
fn from_vec3() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let vec3 = vec![vec![vec![1, 2]; 2]; 2];
    let pyarray = PyArray::from_vec3(gil.python(), &np, &vec3).unwrap();
    assert_eq!(
        pyarray.as_array().unwrap(),
        array![[[1, 2], [1, 2]], [[1, 2], [1, 2]]].into_dyn()
    );
}

#[test]
fn from_small_array() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let array: [i32; 5] = [1, 2, 3, 4, 5];
    let pyarray = array.into_pyarray(gil.python(), &np);
    assert_eq!(pyarray.as_slice().unwrap(), &[1, 2, 3, 4, 5]);
}

#[test]
fn from_eval() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let dict = PyDict::new(gil.python());
    dict.set_item("np", np.as_pymodule()).unwrap();
    let pyarray: &PyArray<i32> = gil
        .python()
        .eval("np.array([1, 2, 3], dtype='int32')", Some(&dict), None)
        .unwrap()
        .extract()
        .unwrap();
    assert_eq!(pyarray.as_slice().unwrap(), &[1, 2, 3]);
}

#[test]
fn from_eval_fail() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let dict = PyDict::new(gil.python());
    dict.set_item("np", np.as_pymodule()).unwrap();
    let converted: Result<&PyArray<i32>, _> = gil
        .python()
        .eval("np.array([1, 2, 3], dtype='float64')", Some(&dict), None)
        .unwrap()
        .extract();
    assert!(converted.is_err());
}

<commit_msg>Fix tests<commit_after>extern crate ndarray;
extern crate numpy;
extern crate pyo3;

use ndarray::*;
use numpy::*;
use pyo3::prelude::*;

#[test]
fn new() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let n = 3;
    let m = 5;
    let arr = PyArray::<f64>::new(gil.python(), &np, &[n, m]);
    assert!(arr.ndim() == 2);
    assert!(arr.dims() == [n, m]);
    assert!(arr.strides() == [m as isize * 8, 8]);
}

#[test]
fn zeros() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let n = 3;
    let m = 5;
    let arr = PyArray::<f64>::zeros(gil.python(), &np, &[n, m], false);
    assert!(arr.ndim() == 2);
    assert!(arr.dims() == [n, m]);
    assert!(arr.strides() == [m as isize * 8, 8]);

    let arr = PyArray::<f64>::zeros(gil.python(), &np, &[n, m], true);
    assert!(arr.ndim() == 2);
    assert!(arr.dims() == [n, m]);
    assert!(arr.strides() == [8, n as isize * 8]);
}

#[test]
fn arange() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let arr = PyArray::<f64>::arange(gil.python(), &np, 0.0, 1.0, 0.1);
    println!("ndim = {:?}", arr.ndim());
    println!("dims = {:?}", arr.dims());
    println!("array = {:?}", arr.as_slice().unwrap());
}

#[test]
fn as_array() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let arr = PyArray::<f64>::zeros(gil.python(), &np, &[3, 2, 4], false);
    let a = arr.as_array().unwrap();
    assert_eq!(arr.shape(), a.shape());
    assert_eq!(
        arr.strides().iter().map(|x| x / 8).collect::<Vec<_>>(),
        a.strides()
    );
}

#[test]
fn into_pyarray_vec() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();

    let a = vec![1, 2, 3];
    let arr = a.into_pyarray(gil.python(), &np);
    println!("arr.shape = {:?}", arr.shape());
    println!("arr = {:?}", arr.as_slice().unwrap());
    assert_eq!(arr.shape(), [3]);
}

#[test]
fn into_pyarray_array() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();

    let a = Array3::<f64>::zeros((3, 4, 2));
    let shape = a.shape().iter().cloned().collect::<Vec<_>>();
    let strides = a.strides().iter().map(|d| d * 8).collect::<Vec<_>>();
    println!("a.shape   = {:?}", a.shape());
    println!("a.strides = {:?}", a.strides());
    let pa = a.into_pyarray(gil.python(), &np);
    println!("pa.shape   = {:?}", pa.shape());
    println!("pa.strides = {:?}", pa.strides());
    assert_eq!(pa.shape(), shape.as_slice());
    assert_eq!(pa.strides(), strides.as_slice());
}

#[test]
fn iter_to_pyarray() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let arr = PyArray::from_iter(gil.python(), &np, (0..10).map(|x| x * x));
    assert_eq!(arr.as_slice().unwrap(), &[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]);
}

#[test]
fn is_instance() {
    let gil = pyo3::Python::acquire_gil();
    let py = gil.python();
    let np = PyArrayModule::import(py).unwrap();
    let arr = PyArray::<f64>::new(gil.python(), &np, &[3, 5]);
    assert!(py.is_instance::<PyArray<f64>, _>(&arr).unwrap());
    assert!(!py.is_instance::<pyo3::PyList, _>(&arr).unwrap());
}

#[test]
fn from_vec2() {
    let vec2 = vec![vec![1, 2, 3]; 2];
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let pyarray = PyArray::from_vec2(gil.python(), &np, &vec2).unwrap();
    assert_eq!(
        pyarray.as_array().unwrap(),
        array![[1, 2, 3], [1, 2, 3]].into_dyn()
    );
    assert!(PyArray::from_vec2(gil.python(), &np, &vec![vec![1], vec![2, 3]]).is_err());
}

#[test]
fn from_vec3() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let vec3 = vec![vec![vec![1, 2]; 2]; 2];
    let pyarray = PyArray::from_vec3(gil.python(), &np, &vec3).unwrap();
    assert_eq!(
        pyarray.as_array().unwrap(),
        array![[[1, 2], [1, 2]], [[1, 2], [1, 2]]].into_dyn()
    );
}

#[test]
fn from_small_array() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let array: [i32; 5] = [1, 2, 3, 4, 5];
    let pyarray = array.into_pyarray(gil.python(), &np);
    assert_eq!(pyarray.as_slice().unwrap(), &[1, 2, 3, 4, 5]);
}

#[test]
fn from_eval() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let dict = PyDict::new(gil.python());
    dict.set_item("np", np.as_pymodule()).unwrap();
    let pyarray: &PyArray<i32> = gil
        .python()
        .eval("np.array([1, 2, 3], dtype='int32')", Some(&dict), None)
        .unwrap()
        .extract()
        .unwrap();
    assert_eq!(pyarray.as_slice().unwrap(), &[1, 2, 3]);
}

#[test]
fn from_eval_fail() {
    let gil = pyo3::Python::acquire_gil();
    let np = PyArrayModule::import(gil.python()).unwrap();
    let dict = PyDict::new(gil.python());
    dict.set_item("np", np.as_pymodule()).unwrap();
    let converted: Result<&PyArray<i32>, _> = gil
        .python()
        .eval("np.array([1, 2, 3], dtype='float64')", Some(&dict), None)
        .unwrap()
        .extract();
    assert!(converted.is_err());
}

<|endoftext|>


<commit_before>#![feature(phase)]
#![feature(globs)]
#![feature(macro_rules)]

#[phase(plugin)]
extern crate deuterium_plugin;

#[phase(plugin)]
extern crate deuterium_orm;
extern crate deuterium_orm;
extern crate time;
extern crate deuterium;

extern crate postgres;
extern crate r2d2;
extern crate r2d2_postgres;
extern crate test;

use deuterium::*;
use deuterium_orm::*;
use time::Timespec;

use postgres::Connection;
use r2d2_postgres::PostgresPoolManager;

macro_rules! assert_sql(
    ($query:expr, $s:expr) => (
        assert_eq!($query.to_final_sql().as_slice(), $s)
    )
)

#[deriving(Clone, Show, PartialEq, FromPrimitive)]
pub enum Side {
    DarkSide,
    LightSide,
}

deuterium_enum!(Side)

deuterium_model! jedi {
    #[primary_key(id)]
    #[before_create(created_at)]
    #[before_save(updated_at)]
    pub struct Jedi {
        id: i32,
        name: String,
        force_level: i32,
        side: Side,
        created_at: Timespec,
        updated_at: Timespec
    }
}

impl Jedi {
    pub fn ordered() -> SelectQuery<(), LimitMany, Jedi> {
        Jedi::table().select_all().order_by(&Jedi::created_at_f())
    }
}

fn created_at(token: &mut Jedi) {
    token.set_created_at(::time::get_time());
}

fn updated_at(token: &mut Jedi) {
    token.set_updated_at(::time::get_time());
}

fn setup_tables(cn: &Connection) {
   cn.batch_execute(r#"
        DROP TABLE IF EXISTS jedi CASCADE;
        CREATE TABLE jedi (
            id          serial PRIMARY KEY,
            name        varchar(40) NOT NULL,
            force_level integer NOT NULL,
            side        SMALLINT NOT NULL,
            created_at  timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
            updated_at  timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL
        );

        INSERT INTO jedi (name, force_level, side) VALUES
            ('Luke Skywalker', 100, 1),
            ('Mace Windu', 90, 1),
            ('Obi-Wan Kenoby', 99, 1),
            ('Kit Fisto', 70, 1),
            ('Count Dooku', 99, 0),
            ('Darth Maul', 70, 0),
            ('Anakin Skywalker', 100, 0);

    "#).unwrap();
}

fn setup_pg() -> adapter::postgres::PostgresPool {
    let manager = PostgresPoolManager::new("postgres://panferov@localhost/jedi", ::postgres::SslMode::None);
    let config = r2d2::Config {
        pool_size: 5,
        test_on_check_out: true,
        ..std::default::Default::default()
    };

    let handler = r2d2::NoopErrorHandler;
    r2d2::Pool::new(config, manager, handler).unwrap()
}

#[test]
fn select() {
    let pool = setup_pg();
    let cn = pool.get().unwrap();

    setup_tables(&*cn);

    assert_eq!((query_models!(
        &Jedi::ordered().where_(Jedi::side_f().is(Side::LightSide)),
        &*cn, &[]
    )).len(), 4);

    let anakin = (query_model!(
        &Jedi::ordered().where_(Jedi::name_f().is("Anakin Skywalker".to_string())).first(),
        &*cn, &[]
    )).unwrap();

    assert_eq!(anakin.get_force_level(), &100)
    assert_eq!(anakin.get_side(), &Side::DarkSide);
}

#[test]
fn insert() {
    let pool = setup_pg();
    let cn = pool.get().unwrap();
    setup_tables(&*cn);

    let mut jedi = Jedi::empty();
    jedi.set_name("Pants Olmos".to_string());
    jedi.set_force_level(10);
    jedi.set_side(Side::DarkSide);

    assert_eq!(exec_pg!(&jedi.create_query(), &*cn, &[]), 1);

    let olmos = (query_model!(
        &Jedi::table().select_all().where_(Jedi::name_f().is("Pants Olmos".to_string())).first(),
        &*cn, &[]
    )).unwrap();

    assert_eq!(olmos.get_name(), &"Pants Olmos".to_string());
    assert_eq!(olmos.get_force_level(), &10);
    assert_eq!(olmos.get_side(), &Side::DarkSide);
}

#[test]
fn update() {
    let pool = setup_pg();
    let cn = pool.get().unwrap();
    setup_tables(&*cn);

    let mut anakin = (query_model!(
        &Jedi::table().select_all().where_(Jedi::name_f().is("Anakin Skywalker".to_string())).first(),
        &*cn, &[]
    )).unwrap();

    assert_eq!(anakin.get_side(), &Side::DarkSide);

    anakin.set_side(Side::LightSide);
    assert_eq!(exec_pg!(&anakin.update_query(), &*cn, &[]), 1);
}

#[test]
fn delete() {
    let pool = setup_pg();
    let cn = pool.get().unwrap();
    setup_tables(&*cn);

    let mut anakin = (query_model!(
        &Jedi::table().select_all().where_(Jedi::name_f().is("Anakin Skywalker".to_string())).first(),
        &*cn, &[]
    )).unwrap();

    assert_eq!(exec_pg!(&anakin.delete_query(), &*cn, &[]), 1);
}<commit_msg>refactor: refactor hardcoded connection string to os env value<commit_after>#![feature(phase)]
#![feature(globs)]
#![feature(macro_rules)]

#[phase(plugin)]
extern crate deuterium_plugin;

#[phase(plugin)]
extern crate deuterium_orm;
extern crate deuterium_orm;
extern crate time;
extern crate deuterium;

extern crate postgres;
extern crate r2d2;
extern crate r2d2_postgres;
extern crate test;

use std::os;
use deuterium::*;
use deuterium_orm::*;
use time::Timespec;

use postgres::Connection;
use r2d2_postgres::PostgresPoolManager;

macro_rules! assert_sql(
    ($query:expr, $s:expr) => (
        assert_eq!($query.to_final_sql().as_slice(), $s)
    )
)

#[deriving(Clone, Show, PartialEq, FromPrimitive)]
pub enum Side {
    DarkSide,
    LightSide,
}

deuterium_enum!(Side)

deuterium_model! jedi {
    #[primary_key(id)]
    #[before_create(created_at)]
    #[before_save(updated_at)]
    pub struct Jedi {
        id: i32,
        name: String,
        force_level: i32,
        side: Side,
        created_at: Timespec,
        updated_at: Timespec
    }
}

impl Jedi {
    pub fn ordered() -> SelectQuery<(), LimitMany, Jedi> {
        Jedi::table().select_all().order_by(&Jedi::created_at_f())
    }
}

fn created_at(token: &mut Jedi) {
    token.set_created_at(::time::get_time());
}

fn updated_at(token: &mut Jedi) {
    token.set_updated_at(::time::get_time());
}

fn setup_tables(cn: &Connection) {
   cn.batch_execute(r#"
        DROP TABLE IF EXISTS jedi CASCADE;
        CREATE TABLE jedi (
            id          serial PRIMARY KEY,
            name        varchar(40) NOT NULL,
            force_level integer NOT NULL,
            side        SMALLINT NOT NULL,
            created_at  timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
            updated_at  timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL
        );

        INSERT INTO jedi (name, force_level, side) VALUES
            ('Luke Skywalker', 100, 1),
            ('Mace Windu', 90, 1),
            ('Obi-Wan Kenoby', 99, 1),
            ('Kit Fisto', 70, 1),
            ('Count Dooku', 99, 0),
            ('Darth Maul', 70, 0),
            ('Anakin Skywalker', 100, 0);

    "#).unwrap();
}

fn setup_pg() -> adapter::postgres::PostgresPool {

    let connection_uri = match os::getenv("POSTGRES_CONNECTION") {
        Some(val) => val,
        None => "postgres://localhost/jedi".to_string()
    };

    let manager = PostgresPoolManager::new(connection_uri.as_slice(), ::postgres::SslMode::None);
    let config = r2d2::Config {
        pool_size: 5,
        test_on_check_out: true,
        ..std::default::Default::default()
    };

    let handler = r2d2::NoopErrorHandler;
    r2d2::Pool::new(config, manager, handler).unwrap()
}

#[test]
fn select() {
    let pool = setup_pg();
    let cn = pool.get().unwrap();

    setup_tables(&*cn);

    assert_eq!((query_models!(
        &Jedi::ordered().where_(Jedi::side_f().is(Side::LightSide)),
        &*cn, &[]
    )).len(), 4);

    let anakin = (query_model!(
        &Jedi::ordered().where_(Jedi::name_f().is("Anakin Skywalker".to_string())).first(),
        &*cn, &[]
    )).unwrap();

    assert_eq!(anakin.get_force_level(), &100)
    assert_eq!(anakin.get_side(), &Side::DarkSide);
}

#[test]
fn insert() {
    let pool = setup_pg();
    let cn = pool.get().unwrap();
    setup_tables(&*cn);

    let mut jedi = Jedi::empty();
    jedi.set_name("Pants Olmos".to_string());
    jedi.set_force_level(10);
    jedi.set_side(Side::DarkSide);

    assert_eq!(exec_pg!(&jedi.create_query(), &*cn, &[]), 1);

    let olmos = (query_model!(
        &Jedi::table().select_all().where_(Jedi::name_f().is("Pants Olmos".to_string())).first(),
        &*cn, &[]
    )).unwrap();

    assert_eq!(olmos.get_name(), &"Pants Olmos".to_string());
    assert_eq!(olmos.get_force_level(), &10);
    assert_eq!(olmos.get_side(), &Side::DarkSide);
}

#[test]
fn update() {
    let pool = setup_pg();
    let cn = pool.get().unwrap();
    setup_tables(&*cn);

    let mut anakin = (query_model!(
        &Jedi::table().select_all().where_(Jedi::name_f().is("Anakin Skywalker".to_string())).first(),
        &*cn, &[]
    )).unwrap();

    assert_eq!(anakin.get_side(), &Side::DarkSide);

    anakin.set_side(Side::LightSide);
    assert_eq!(exec_pg!(&anakin.update_query(), &*cn, &[]), 1);
}

#[test]
fn delete() {
    let pool = setup_pg();
    let cn = pool.get().unwrap();
    setup_tables(&*cn);

    let mut anakin = (query_model!(
        &Jedi::table().select_all().where_(Jedi::name_f().is("Anakin Skywalker".to_string())).first(),
        &*cn, &[]
    )).unwrap();

    assert_eq!(exec_pg!(&anakin.delete_query(), &*cn, &[]), 1);
}<|endoftext|>


<commit_before>extern crate cbor;
extern crate quickcheck;
extern crate rand;
extern crate rustc_serialize;

use std::collections::HashMap;
use std::fmt::Debug;
use rand::thread_rng;
use rustc_serialize::{Decodable, Encodable};
use quickcheck::{QuickCheck, StdGen, Testable};

use cbor::{Encoder, Decoder, Cbor, CborBytes, CborTagEncode};

fn qc_sized<A: Testable>(f: A, size: u64) {
    QuickCheck::new()
        .gen(StdGen::new(thread_rng(), size as usize))
        .tests(1_00)
        .max_tests(10_000)
        .quickcheck(f);
}

fn round_trip<T>(v: T) -> bool
        where T: Decodable + Encodable + Debug + PartialEq {
    let backv: T = decode(&encode(&v));
    assert_eq!(backv, v);
    true
}

fn encode<T: Encodable>(v: T) -> Vec<u8> {
    let mut enc = Encoder::from_memory();
    enc.encode(&[v]).unwrap();
    enc.as_bytes().to_vec()
}

fn decode<T: Decodable>(bytes: &[u8]) -> T {
    Decoder::from_bytes(bytes).decode().next().unwrap().unwrap()
}

#[allow(dead_code)]
fn readone(bytes: &[u8]) -> Cbor {
    Decoder::from_bytes(bytes).items().next().unwrap().unwrap()
}

macro_rules! round_trip_num {
    ($name:ident, $ty:ident) => (
        #[test]
        fn $name() {
            #![allow(trivial_numeric_casts)]
            fn prop(n: $ty) -> bool { round_trip(n) }
            qc_sized(prop as fn($ty) -> bool, ::std::$ty::MAX as u64)
        }
    );
}

round_trip_num!(roundtrip_prop_usize, usize);
round_trip_num!(roundtrip_prop_u64, u64);
round_trip_num!(roundtrip_prop_u32, u32);
round_trip_num!(roundtrip_prop_u16, u16);
round_trip_num!(roundtrip_prop_u8, u8);
round_trip_num!(roundtrip_prop_isize, isize);
round_trip_num!(roundtrip_prop_i64, i64);
round_trip_num!(roundtrip_prop_i32, i32);
round_trip_num!(roundtrip_prop_i16, i16);
round_trip_num!(roundtrip_prop_i8, i8);

#[test]
fn roundtrip_prop_f32() {
    fn prop(n: f32) -> bool { round_trip(n) }
    qc_sized(prop as fn(f32) -> bool, ::std::u64::MAX)
}

#[test]
fn roundtrip_prop_f64() {
    fn prop(n: f64) -> bool { round_trip(n) }
    qc_sized(prop as fn(f64) -> bool, ::std::u64::MAX)
}

#[test]
fn roundtrip_f32_min() {
    round_trip(::std::f32::MIN);
}

#[test]
fn roundtrip_f32_max() {
    round_trip(::std::f32::MAX);
}

#[test]
fn roundtrip_f64_min() {
    round_trip(::std::f64::MIN);
}

#[test]
fn roundtrip_f64_max() {
    round_trip(::std::f64::MAX);
}

macro_rules! round_trip_any {
    ($name:ident, $ty:ty) => (
        #[test]
        fn $name() {
            fn prop(n: $ty) -> bool { round_trip(n) }
            QuickCheck::new().quickcheck(prop as fn($ty) -> bool)
        }
    );
}

round_trip_any!(roundtrip_prop_nil, ());
round_trip_any!(roundtrip_prop_bool, bool);
round_trip_any!(roundtrip_prop_tuple, (i32, u32));
round_trip_any!(roundtrip_prop_tuple_in_tuple, (i32, (u32, i32), u32));
round_trip_any!(roundtrip_prop_vec, Vec<i32>);
round_trip_any!(roundtrip_prop_vec_in_vec, Vec<Vec<i32>>);
round_trip_any!(roundtrip_prop_map, HashMap<String, i32>);
round_trip_any!(roundtrip_prop_vec_in_map, HashMap<String, Vec<i32>>);
// round_trip_any!(roundtrip_prop_map_in_map,
                // HashMap<String, HashMap<String, i32>>);

#[test]
fn roundtrip_prop_byte_string() {
    fn prop(n: Vec<u8>) -> bool { round_trip(CborBytes(n)) }
    QuickCheck::new().quickcheck(prop as fn(Vec<u8>) -> bool)
}

#[test]
fn roundtrip_enum() {
    #[derive(Debug, PartialEq, RustcDecodable, RustcEncodable)]
    enum Color { Red, Blue(String, i32), Green { s: String, n: i32 } }

    round_trip(Color::Red);
    round_trip(Color::Blue("hi".to_string(), 5));
    round_trip(Color::Green { s: "hi".to_string(), n: 5 });
}

#[test]
fn roundtrip_struct() {
    #[derive(Debug, PartialEq, RustcDecodable, RustcEncodable)]
    struct Vowels { s: String, n: u32 }

    round_trip(Vowels { s: "cwm".to_string(), n: 1 });
}

#[test]
#[should_panic]
fn invalid_map_key() {
    let mut map = HashMap::new();
    map.insert(5, 5);
    encode(map);
}

#[test]
fn roundtrip_prop_tag() {
    use rustc_serialize::{Decoder, Encoder};

    #[derive(Debug, PartialEq)]
    struct MyTag {
        num: u64,
        data: Vec<i32>,
    }
    impl Decodable for MyTag {
        fn decode<D: Decoder>(d: &mut D) -> Result<MyTag, D::Error> {
            let tag = try!(d.read_u64());
            Ok(MyTag { num: tag, data: try!(Decodable::decode(d)) })
        }
    }
    impl Encodable for MyTag {
        fn encode<E: Encoder>(&self, e: &mut E) -> Result<(), E::Error> {
            CborTagEncode::new(self.num, &self.data).encode(e)
        }
    }
    fn prop(num: u64, data: Vec<i32>) -> bool {
        round_trip(MyTag { num: num, data: data })
    }
    QuickCheck::new().quickcheck(prop as fn(u64, Vec<i32>) -> bool)
}

#[test]
fn roundtrip_tag_fancier_data() {
    use rustc_serialize::{Decoder, Encoder};

    #[derive(Debug, PartialEq, RustcDecodable, RustcEncodable)]
    struct DataName(Vec<u8>);

    #[derive(Debug, PartialEq)]
    struct CustomData {
        name: DataName,
        value: Vec<u8>,
    }

    impl Decodable for CustomData {
        fn decode<D: Decoder>(d: &mut D) -> Result<CustomData, D::Error> {
            let _ = try!(d.read_u64());
            let (name, value) = try!(Decodable::decode(d));
            Ok(CustomData { name: name, value: value })
        }
    }
    impl Encodable for CustomData {
        fn encode<E: Encoder>(&self, e: &mut E) -> Result<(), E::Error> {
            CborTagEncode::new(100_000, &(&self.name, &self.value)).encode(e)
        }
    }
    assert!(round_trip(CustomData {
        name: DataName(b"hello".to_vec()),
        value: vec![1, 2, 3, 4, 5],
    }));
}
<commit_msg>Add a test case that exibits over-capacity out-of-memory error<commit_after>extern crate cbor;
extern crate quickcheck;
extern crate rand;
extern crate rustc_serialize;

use std::collections::HashMap;
use std::fmt::Debug;
use rand::thread_rng;
use rustc_serialize::{Decodable, Encodable};
use quickcheck::{QuickCheck, StdGen, Testable};

use cbor::{Encoder, Decoder, Cbor, CborBytes, CborTagEncode};

fn qc_sized<A: Testable>(f: A, size: u64) {
    QuickCheck::new()
        .gen(StdGen::new(thread_rng(), size as usize))
        .tests(1_00)
        .max_tests(10_000)
        .quickcheck(f);
}

fn round_trip<T>(v: T) -> bool
        where T: Decodable + Encodable + Debug + PartialEq {
    let backv: T = decode(&encode(&v));
    assert_eq!(backv, v);
    true
}

fn encode<T: Encodable>(v: T) -> Vec<u8> {
    let mut enc = Encoder::from_memory();
    enc.encode(&[v]).unwrap();
    enc.as_bytes().to_vec()
}

fn decode<T: Decodable>(bytes: &[u8]) -> T {
    Decoder::from_bytes(bytes).decode().next().unwrap().unwrap()
}

#[allow(dead_code)]
fn readone(bytes: &[u8]) -> Cbor {
    Decoder::from_bytes(bytes).items().next().unwrap().unwrap()
}

macro_rules! round_trip_num {
    ($name:ident, $ty:ident) => (
        #[test]
        fn $name() {
            #![allow(trivial_numeric_casts)]
            fn prop(n: $ty) -> bool { round_trip(n) }
            qc_sized(prop as fn($ty) -> bool, ::std::$ty::MAX as u64)
        }
    );
}

round_trip_num!(roundtrip_prop_usize, usize);
round_trip_num!(roundtrip_prop_u64, u64);
round_trip_num!(roundtrip_prop_u32, u32);
round_trip_num!(roundtrip_prop_u16, u16);
round_trip_num!(roundtrip_prop_u8, u8);
round_trip_num!(roundtrip_prop_isize, isize);
round_trip_num!(roundtrip_prop_i64, i64);
round_trip_num!(roundtrip_prop_i32, i32);
round_trip_num!(roundtrip_prop_i16, i16);
round_trip_num!(roundtrip_prop_i8, i8);

#[test]
fn roundtrip_prop_f32() {
    fn prop(n: f32) -> bool { round_trip(n) }
    qc_sized(prop as fn(f32) -> bool, ::std::u64::MAX)
}

#[test]
fn roundtrip_prop_f64() {
    fn prop(n: f64) -> bool { round_trip(n) }
    qc_sized(prop as fn(f64) -> bool, ::std::u64::MAX)
}

#[test]
fn roundtrip_f32_min() {
    round_trip(::std::f32::MIN);
}

#[test]
fn roundtrip_f32_max() {
    round_trip(::std::f32::MAX);
}

#[test]
fn roundtrip_f64_min() {
    round_trip(::std::f64::MIN);
}

#[test]
fn roundtrip_f64_max() {
    round_trip(::std::f64::MAX);
}

macro_rules! round_trip_any {
    ($name:ident, $ty:ty) => (
        #[test]
        fn $name() {
            fn prop(n: $ty) -> bool { round_trip(n) }
            QuickCheck::new().quickcheck(prop as fn($ty) -> bool)
        }
    );
}

round_trip_any!(roundtrip_prop_nil, ());
round_trip_any!(roundtrip_prop_bool, bool);
round_trip_any!(roundtrip_prop_tuple, (i32, u32));
round_trip_any!(roundtrip_prop_tuple_in_tuple, (i32, (u32, i32), u32));
round_trip_any!(roundtrip_prop_vec, Vec<i32>);
round_trip_any!(roundtrip_prop_vec_in_vec, Vec<Vec<i32>>);
round_trip_any!(roundtrip_prop_map, HashMap<String, i32>);
round_trip_any!(roundtrip_prop_vec_in_map, HashMap<String, Vec<i32>>);
// round_trip_any!(roundtrip_prop_map_in_map,
                // HashMap<String, HashMap<String, i32>>);

#[test]
fn roundtrip_prop_byte_string() {
    fn prop(n: Vec<u8>) -> bool { round_trip(CborBytes(n)) }
    QuickCheck::new().quickcheck(prop as fn(Vec<u8>) -> bool)
}

#[test]
fn roundtrip_enum() {
    #[derive(Debug, PartialEq, RustcDecodable, RustcEncodable)]
    enum Color { Red, Blue(String, i32), Green { s: String, n: i32 } }

    round_trip(Color::Red);
    round_trip(Color::Blue("hi".to_string(), 5));
    round_trip(Color::Green { s: "hi".to_string(), n: 5 });
}

#[test]
fn roundtrip_struct() {
    #[derive(Debug, PartialEq, RustcDecodable, RustcEncodable)]
    struct Vowels { s: String, n: u32 }

    round_trip(Vowels { s: "cwm".to_string(), n: 1 });
}

#[test]
#[should_panic]
fn invalid_map_key() {
    let mut map = HashMap::new();
    map.insert(5, 5);
    encode(map);
}

#[test]
fn roundtrip_prop_tag() {
    use rustc_serialize::{Decoder, Encoder};

    #[derive(Debug, PartialEq)]
    struct MyTag {
        num: u64,
        data: Vec<i32>,
    }
    impl Decodable for MyTag {
        fn decode<D: Decoder>(d: &mut D) -> Result<MyTag, D::Error> {
            let tag = try!(d.read_u64());
            Ok(MyTag { num: tag, data: try!(Decodable::decode(d)) })
        }
    }
    impl Encodable for MyTag {
        fn encode<E: Encoder>(&self, e: &mut E) -> Result<(), E::Error> {
            CborTagEncode::new(self.num, &self.data).encode(e)
        }
    }
    fn prop(num: u64, data: Vec<i32>) -> bool {
        round_trip(MyTag { num: num, data: data })
    }
    QuickCheck::new().quickcheck(prop as fn(u64, Vec<i32>) -> bool)
}

#[test]
fn roundtrip_tag_fancier_data() {
    use rustc_serialize::{Decoder, Encoder};

    #[derive(Debug, PartialEq, RustcDecodable, RustcEncodable)]
    struct DataName(Vec<u8>);

    #[derive(Debug, PartialEq)]
    struct CustomData {
        name: DataName,
        value: Vec<u8>,
    }

    impl Decodable for CustomData {
        fn decode<D: Decoder>(d: &mut D) -> Result<CustomData, D::Error> {
            let _ = try!(d.read_u64());
            let (name, value) = try!(Decodable::decode(d));
            Ok(CustomData { name: name, value: value })
        }
    }
    impl Encodable for CustomData {
        fn encode<E: Encoder>(&self, e: &mut E) -> Result<(), E::Error> {
            CborTagEncode::new(100_000, &(&self.name, &self.value)).encode(e)
        }
    }
    assert!(round_trip(CustomData {
        name: DataName(b"hello".to_vec()),
        value: vec![1, 2, 3, 4, 5],
    }));
}

#[test]
fn test_oom() {
   let bad = vec![155u8, 0x00, 0xFF, 0xFF, 0xFF, 0x00, 0xFF, 0xFF, 0xFF];
   let _: Vec<u32> = decode(&bad[..]);
}
<|endoftext|>


<commit_before>use graph_algorithms::Graph;
use graph_algorithms::node_vec::NodeVec;
use graph::{BasicBlockIndex, FuncGraph};
use env::{Environment, Point};
use std::collections::BTreeMap;
use std::cmp;
use std::fmt;

/// A region is fully characterized by a set of tails.  The entry
/// point is the innerloop loop contained by one of the tails. The
/// region consists of all blocks that are dominated by the entry
/// point **and** which dominate one of the tails.
///
/// To be self-consistent:
/// - All tails must be dominated by the entry point.
/// - There must not exist an edge U->V where U is not in the region but
///   V is, unless V is the entry point.
#[derive(Clone, PartialEq, Eq)]
pub struct Region {
    entry: BasicBlockIndex,
    leaves: BTreeMap<BasicBlockIndex, usize>,
}

impl Region {
    pub fn with_point(point: Point) -> Self {
        // The minimal region containing point just has `point` as a tail.
        Self::new(point.block, Some(point))
    }

    pub fn new<I>(entry: BasicBlockIndex, leaves: I) -> Self
        where I: IntoIterator<Item = Point>
    {
        Region {
            entry: entry,
            leaves: leaves.into_iter()
                          .map(|p| (p.block, p.action))
                          .collect(),
        }
    }

    pub fn leaves(&self) -> Leaves {
        Leaves {
            iter: Box::new(self.leaves
                               .iter()
                               .map(|(&block, &action)| {
                                   Point {
                                       block: block,
                                       action: action,
                                   }
                               })),
        }
    }

    pub fn contains(&self, env: &Environment, point: Point) -> bool {
        log!("contains(self={:?}, point={:?})", self, point);

        // no leaves is the empty region, which contains no points
        if self.leaves.is_empty() {
            return false;
        }

        // point must be dominated by the entry
        if !env.dominators.is_dominated_by(point.block, self.entry) {
            log!("contains: {:?} not dominated by entry {:?}",
                 point.block,
                 self.entry);
            return false;
        }

        // point dominate one of the leaves; if it is in a tail block,
        // compare the actions, otherwise consult the dominator graph
        if let Some(&action) = self.leaves.get(&point.block) {
            log!("contains: {:?} is a tail upto {:?}", point.block, action);
            return point.action <= action;
        }
        log!("contains: comparing leaves");
        self.leaves
            .keys()
            .any(|&tail| env.dominators.is_dominated_by(tail, point.block))
    }

    pub fn add_point(&mut self, env: &Environment, point: Point) -> bool {
        log!("add_point(self={:?}, point={:?})", self, point);
        assert!(point.action < env.end_action(point.block));

        if self.contains(env, point) {
            return false;
        }

        // Update entry point to be the mutual dominator
        self.entry = env.dominators.mutual_dominator_node(self.entry, point.block);

        // Compute the minimal consistent region beginning at
        // `new_entry` that contains all the leaves and the new point.
        let mut rev_dfs = RevDfs::new(env, self.entry);
        for tail in self.leaves() {
            rev_dfs.walk(tail);
        }
        rev_dfs.walk(point);

        // Distill that into just the leaves and update our internal map.
        self.leaves.clear();
        let mut stack = vec![self.entry];
        while let Some(node) = stack.pop() {
            let upto_action = rev_dfs.contained[node];
            let end_action = env.end_action(node);

            log!("node {:?} contained up to action {:?} out of {:?}",
                 node,
                 upto_action,
                 end_action);

            assert!(upto_action <= end_action);
            if upto_action == end_action {
                // Block is fully contained; visit its children,
                // unless this is a leaf of the tree, in which case we
                // have to record is as one of our leaves.
                let children = env.dominator_tree.children(node);
                if children.is_empty() {
                    self.leaves.insert(node, upto_action - 1);
                } else {
                    stack.extend(children);
                }
                continue;
            }

            if upto_action == 0 {
                // Block is not contained at all. Stop walking.
                continue;
            }

            // Block is partially contained. Record as a tail but stop
            // walking.
            self.leaves.insert(node, upto_action - 1);
        }

        log!("add_point: leaves={:?}", self.leaves);
        assert!(self.contains(env, point));
        true
    }
}

struct RevDfs<'env, 'func: 'env, 'arena: 'func> {
    env: &'env Environment<'func, 'arena>,

    // for each block, stores the number of actions that are contained
    // in the set; hence a value of `0` means that a block is not
    // contained at all; a value of `1` means the first action is
    // contained but not any others; and a value of `N` where `N` is
    // the number of actions means the block is fully contained.
    contained: NodeVec<FuncGraph<'arena>, usize>,
    stack: Vec<BasicBlockIndex>,
    entry: BasicBlockIndex,
}

impl<'env, 'func, 'arena> RevDfs<'env, 'func, 'arena> {
    fn new(env: &'env Environment<'func, 'arena>, entry: BasicBlockIndex) -> Self {
        RevDfs {
            env: env,
            contained: NodeVec::from_default(env.graph),
            stack: vec![],
            entry: entry,
        }
    }

    fn walk(&mut self, mut start: Point) {
        log!("walk({:?})", start);
        assert!(self.env.dominators.is_dominated_by(start.block, self.entry));
        assert!(self.stack.is_empty());

        // convert `start` to be exclusive (one past the included point)
        start.action += 1;

        self.push(start);

        // Visit reachable predecessors, unless this is the region entry
        while let Some(p) = self.pop() {
            log!("walk: walking preds of {:?}", p);
            if p == self.entry {
                continue;
            }
            for pred_block in self.env.graph.predecessors(p) {
                self.push(self.env.end_point(pred_block));
            }
        }
    }

    fn push(&mut self, p: Point) {
        log!("push({:?})", p);
        if self.contained[p.block] > 0 {
            log!("push: already visited {:?} up to action {:?} out of {:?}, \
                  adjusting to {:?}",
                 p.block,
                 self.contained[p.block],
                 self.env.end_action(p.block),
                 cmp::max(p.action, self.contained[p.block]));
            self.contained[p.block] = cmp::max(p.action, self.contained[p.block]);
        } else {
            log!("push: enstacking {:?} up to action {:?}", p.block, p.action);
            self.contained[p.block] = p.action;
            self.stack.push(p.block); // still have to visit preds too
        }
    }

    fn pop(&mut self) -> Option<BasicBlockIndex> {
        self.stack.pop()
    }
}

pub struct Leaves<'iter> {
    iter: Box<Iterator<Item = Point> + 'iter>,
}

impl<'iter> Iterator for Leaves<'iter> {
    type Item = Point;

    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next()
    }
}

impl fmt::Debug for Region {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        try!(write!(fmt, "{{{:?} -> ", self.entry));
        for leaf in self.leaves() {
            try!(write!(fmt, "{:?}", leaf));
        }
        write!(fmt, "}}")
    }
}
<commit_msg>update outdated comment<commit_after>use graph_algorithms::Graph;
use graph_algorithms::node_vec::NodeVec;
use graph::{BasicBlockIndex, FuncGraph};
use env::{Environment, Point};
use std::collections::BTreeMap;
use std::cmp;
use std::fmt;

/// A region is characterized by an entry point and a set of leaves.
/// The region contains all blocks B where:
/// - `entry dom B`
/// - and `exists leaf in leaves. B dom leaf`
#[derive(Clone, PartialEq, Eq)]
pub struct Region {
    entry: BasicBlockIndex,
    leaves: BTreeMap<BasicBlockIndex, usize>,
}

impl Region {
    pub fn with_point(point: Point) -> Self {
        // The minimal region containing point just has `point` as a tail.
        Self::new(point.block, Some(point))
    }

    pub fn new<I>(entry: BasicBlockIndex, leaves: I) -> Self
        where I: IntoIterator<Item = Point>
    {
        Region {
            entry: entry,
            leaves: leaves.into_iter()
                          .map(|p| (p.block, p.action))
                          .collect(),
        }
    }

    pub fn leaves(&self) -> Leaves {
        Leaves {
            iter: Box::new(self.leaves
                               .iter()
                               .map(|(&block, &action)| {
                                   Point {
                                       block: block,
                                       action: action,
                                   }
                               })),
        }
    }

    pub fn contains(&self, env: &Environment, point: Point) -> bool {
        log!("contains(self={:?}, point={:?})", self, point);

        // no leaves is the empty region, which contains no points
        if self.leaves.is_empty() {
            return false;
        }

        // point must be dominated by the entry
        if !env.dominators.is_dominated_by(point.block, self.entry) {
            log!("contains: {:?} not dominated by entry {:?}",
                 point.block,
                 self.entry);
            return false;
        }

        // point dominate one of the leaves; if it is in a tail block,
        // compare the actions, otherwise consult the dominator graph
        if let Some(&action) = self.leaves.get(&point.block) {
            log!("contains: {:?} is a tail upto {:?}", point.block, action);
            return point.action <= action;
        }
        log!("contains: comparing leaves");
        self.leaves
            .keys()
            .any(|&tail| env.dominators.is_dominated_by(tail, point.block))
    }

    pub fn add_point(&mut self, env: &Environment, point: Point) -> bool {
        log!("add_point(self={:?}, point={:?})", self, point);
        assert!(point.action < env.end_action(point.block));

        if self.contains(env, point) {
            return false;
        }

        // Update entry point to be the mutual dominator
        self.entry = env.dominators.mutual_dominator_node(self.entry, point.block);

        // Compute the minimal consistent region beginning at
        // `new_entry` that contains all the leaves and the new point.
        let mut rev_dfs = RevDfs::new(env, self.entry);
        for tail in self.leaves() {
            rev_dfs.walk(tail);
        }
        rev_dfs.walk(point);

        // Distill that into just the leaves and update our internal map.
        self.leaves.clear();
        let mut stack = vec![self.entry];
        while let Some(node) = stack.pop() {
            let upto_action = rev_dfs.contained[node];
            let end_action = env.end_action(node);

            log!("node {:?} contained up to action {:?} out of {:?}",
                 node,
                 upto_action,
                 end_action);

            assert!(upto_action <= end_action);
            if upto_action == end_action {
                // Block is fully contained; visit its children,
                // unless this is a leaf of the tree, in which case we
                // have to record is as one of our leaves.
                let children = env.dominator_tree.children(node);
                if children.is_empty() {
                    self.leaves.insert(node, upto_action - 1);
                } else {
                    stack.extend(children);
                }
                continue;
            }

            if upto_action == 0 {
                // Block is not contained at all. Stop walking.
                continue;
            }

            // Block is partially contained. Record as a tail but stop
            // walking.
            self.leaves.insert(node, upto_action - 1);
        }

        log!("add_point: leaves={:?}", self.leaves);
        assert!(self.contains(env, point));
        true
    }
}

struct RevDfs<'env, 'func: 'env, 'arena: 'func> {
    env: &'env Environment<'func, 'arena>,

    // for each block, stores the number of actions that are contained
    // in the set; hence a value of `0` means that a block is not
    // contained at all; a value of `1` means the first action is
    // contained but not any others; and a value of `N` where `N` is
    // the number of actions means the block is fully contained.
    contained: NodeVec<FuncGraph<'arena>, usize>,
    stack: Vec<BasicBlockIndex>,
    entry: BasicBlockIndex,
}

impl<'env, 'func, 'arena> RevDfs<'env, 'func, 'arena> {
    fn new(env: &'env Environment<'func, 'arena>, entry: BasicBlockIndex) -> Self {
        RevDfs {
            env: env,
            contained: NodeVec::from_default(env.graph),
            stack: vec![],
            entry: entry,
        }
    }

    fn walk(&mut self, mut start: Point) {
        log!("walk({:?})", start);
        assert!(self.env.dominators.is_dominated_by(start.block, self.entry));
        assert!(self.stack.is_empty());

        // convert `start` to be exclusive (one past the included point)
        start.action += 1;

        self.push(start);

        // Visit reachable predecessors, unless this is the region entry
        while let Some(p) = self.pop() {
            log!("walk: walking preds of {:?}", p);
            if p == self.entry {
                continue;
            }
            for pred_block in self.env.graph.predecessors(p) {
                self.push(self.env.end_point(pred_block));
            }
        }
    }

    fn push(&mut self, p: Point) {
        log!("push({:?})", p);
        if self.contained[p.block] > 0 {
            log!("push: already visited {:?} up to action {:?} out of {:?}, \
                  adjusting to {:?}",
                 p.block,
                 self.contained[p.block],
                 self.env.end_action(p.block),
                 cmp::max(p.action, self.contained[p.block]));
            self.contained[p.block] = cmp::max(p.action, self.contained[p.block]);
        } else {
            log!("push: enstacking {:?} up to action {:?}", p.block, p.action);
            self.contained[p.block] = p.action;
            self.stack.push(p.block); // still have to visit preds too
        }
    }

    fn pop(&mut self) -> Option<BasicBlockIndex> {
        self.stack.pop()
    }
}

pub struct Leaves<'iter> {
    iter: Box<Iterator<Item = Point> + 'iter>,
}

impl<'iter> Iterator for Leaves<'iter> {
    type Item = Point;

    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next()
    }
}

impl fmt::Debug for Region {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        try!(write!(fmt, "{{{:?} -> ", self.entry));
        for leaf in self.leaves() {
            try!(write!(fmt, "{:?}", leaf));
        }
        write!(fmt, "}}")
    }
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Defines aggregations over Arrow arrays.

use std::ops::Add;

use crate::array::{Array, GenericStringArray, PrimitiveArray, StringOffsetSizeTrait};
use crate::datatypes::ArrowNumericType;

/// Helper macro to perform min/max of strings
fn min_max_string<T: StringOffsetSizeTrait, F: Fn(&str, &str) -> bool>(
    array: &GenericStringArray<T>,
    cmp: F,
) -> Option<&str> {
    let null_count = array.null_count();

    if null_count == array.len() {
        return None;
    }
    let data = array.data();
    let mut n;
    if null_count == 0 {
        n = array.value(0);
        for i in 1..data.len() {
            let item = array.value(i);
            if cmp(&n, item) {
                n = item;
            }
        }
    } else {
        n = "";
        let mut has_value = false;

        for i in 0..data.len() {
            let item = array.value(i);
            if data.is_valid(i) && (!has_value || cmp(&n, item)) {
                has_value = true;
                n = item;
            }
        }
    }
    Some(n)
}

/// Returns the minimum value in the array, according to the natural order.
pub fn min<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
where
    T: ArrowNumericType,
{
    min_max_helper(array, |a, b| a > b)
}

/// Returns the maximum value in the array, according to the natural order.
pub fn max<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
where
    T: ArrowNumericType,
{
    min_max_helper(array, |a, b| a < b)
}

/// Returns the maximum value in the string array, according to the natural order.
pub fn max_string<T: StringOffsetSizeTrait>(
    array: &GenericStringArray<T>,
) -> Option<&str> {
    min_max_string(array, |a, b| a < b)
}

/// Returns the minimum value in the string array, according to the natural order.
pub fn min_string<T: StringOffsetSizeTrait>(
    array: &GenericStringArray<T>,
) -> Option<&str> {
    min_max_string(array, |a, b| a > b)
}

/// Helper function to perform min/max lambda function on values from a numeric array.
fn min_max_helper<T, F>(array: &PrimitiveArray<T>, cmp: F) -> Option<T::Native>
where
    T: ArrowNumericType,
    F: Fn(&T::Native, &T::Native) -> bool,
{
    let null_count = array.null_count();

    // Includes case array.len() == 0
    if null_count == array.len() {
        return None;
    }

    let data = array.data();
    let m = array.value_slice(0, data.len());
    let mut n;

    if null_count == 0 {
        // optimized path for arrays without null values
        n = m[1..]
            .iter()
            .fold(m[0], |max, item| if cmp(&max, item) { *item } else { max });
    } else {
        n = T::default_value();
        let mut has_value = false;
        for (i, item) in m.iter().enumerate() {
            if data.is_valid(i) && (!has_value || cmp(&n, item)) {
                has_value = true;
                n = *item
            }
        }
    }
    Some(n)
}

/// Returns the sum of values in the array.
///
/// Returns `None` if the array is empty or only contains null values.
#[cfg(not(all(any(target_arch = "x86", target_arch = "x86_64"), feature = "simd")))]
pub fn sum<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
where
    T: ArrowNumericType,
    T::Native: Add<Output = T::Native>,
{
    let null_count = array.null_count();

    if null_count == array.len() {
        return None;
    }

    let data: &[T::Native] = array.value_slice(0, array.len());

    match array.data().null_buffer() {
        None => {
            let sum = data.iter().fold(T::default_value(), |accumulator, value| {
                accumulator + *value
            });

            Some(sum)
        }
        Some(buffer) => {
            let mut sum = T::default_value();
            let data_chunks = data.chunks_exact(64);
            let remainder = data_chunks.remainder();

            let bit_chunks = buffer.bit_chunks(array.offset(), array.len());
            data_chunks
                .zip(bit_chunks.iter())
                .for_each(|(chunk, mask)| {
                    chunk.iter().enumerate().for_each(|(i, value)| {
                        if (mask & (1 << i)) != 0 {
                            sum = sum + *value;
                        }
                    });
                });

            let remainder_bits = bit_chunks.remainder_bits();

            remainder.iter().enumerate().for_each(|(i, value)| {
                if remainder_bits & (1 << i) != 0 {
                    sum = sum + *value;
                }
            });

            Some(sum)
        }
    }
}

/// Returns the sum of values in the array.
///
/// Returns `None` if the array is empty or only contains null values.
#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), feature = "simd"))]
pub fn sum<T: ArrowNumericType>(array: &PrimitiveArray<T>) -> Option<T::Native>
where
    T::Native: Add<Output = T::Native>,
{
    let null_count = array.null_count();

    if null_count == array.len() {
        return None;
    }

    let data: &[T::Native] = array.value_slice(0, array.len());

    let mut vector_sum = T::init(T::default_value());
    let mut remainder_sum = T::default_value();

    match array.data().null_buffer() {
        None => {
            let data_chunks = data.chunks_exact(64);
            let remainder = data_chunks.remainder();

            data_chunks.for_each(|chunk| {
                chunk.chunks_exact(T::lanes()).for_each(|chunk| {
                    let chunk = T::load(&chunk);
                    vector_sum = vector_sum + chunk;
                });
            });

            remainder.iter().for_each(|value| {
                remainder_sum = remainder_sum + *value;
            });
        }
        Some(buffer) => {
            // process data in chunks of 64 elements since we also get 64 bits of validity information at a time
            let data_chunks = data.chunks_exact(64);
            let remainder = data_chunks.remainder();

            let bit_chunks = buffer.bit_chunks(array.offset(), array.len());
            let remainder_bits = bit_chunks.remainder_bits();

            data_chunks.zip(bit_chunks).for_each(|(chunk, mut mask)| {
                // split chunks further into slices corresponding to the vector length
                // the compiler is able to unroll this inner loop and remove bounds checks
                // since the outer chunk size (64) is always a multiple of the number of lanes
                chunk.chunks_exact(T::lanes()).for_each(|chunk| {
                    let zero = T::init(T::default_value());
                    let vecmask = T::mask_from_u64(mask);
                    let chunk = T::load(&chunk);
                    let blended = T::mask_select(vecmask, chunk, zero);

                    vector_sum = vector_sum + blended;

                    mask = mask >> T::lanes();
                });
            });

            remainder.iter().enumerate().for_each(|(i, value)| {
                if remainder_bits & (1 << i) != 0 {
                    remainder_sum = remainder_sum + *value;
                }
            });
        }
    }

    // calculate horizontal sum of accumulator by writing to a temporary
    // this is probably faster than extracting individual lanes
    // the compiler is free to optimize this to something faster
    let tmp = &mut [T::default_value(); 64];
    T::write(vector_sum, &mut tmp[0..T::lanes()]);

    let mut total_sum = T::default_value();
    tmp[0..T::lanes()]
        .iter()
        .for_each(|lane| total_sum = total_sum + *lane);

    total_sum = total_sum + remainder_sum;

    Some(total_sum)
}

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

    #[test]
    fn test_primitive_array_sum() {
        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
        assert_eq!(15, sum(&a).unwrap());
    }

    #[test]
    fn test_primitive_array_float_sum() {
        let a = Float64Array::from(vec![1.1, 2.2, 3.3, 4.4, 5.5]);
        assert!(16.5 - sum(&a).unwrap() < f64::EPSILON);
    }

    #[test]
    fn test_primitive_array_sum_with_nulls() {
        let a = Int32Array::from(vec![None, Some(2), Some(3), None, Some(5)]);
        assert_eq!(10, sum(&a).unwrap());
    }

    #[test]
    fn test_primitive_array_sum_all_nulls() {
        let a = Int32Array::from(vec![None, None, None]);
        assert_eq!(None, sum(&a));
    }

    #[test]
    fn test_buffer_array_min_max() {
        let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
        assert_eq!(5, min(&a).unwrap());
        assert_eq!(9, max(&a).unwrap());
    }

    #[test]
    fn test_buffer_array_min_max_with_nulls() {
        let a = Int32Array::from(vec![Some(5), None, None, Some(8), Some(9)]);
        assert_eq!(5, min(&a).unwrap());
        assert_eq!(9, max(&a).unwrap());
    }

    #[test]
    fn test_buffer_min_max_1() {
        let a = Int32Array::from(vec![None, None, Some(5), Some(2)]);
        assert_eq!(Some(2), min(&a));
        assert_eq!(Some(5), max(&a));
    }

    #[test]
    fn test_string_min_max_with_nulls() {
        let a = StringArray::from(vec![Some("b"), None, None, Some("a"), Some("c")]);
        assert_eq!("a", min_string(&a).unwrap());
        assert_eq!("c", max_string(&a).unwrap());
    }

    #[test]
    fn test_string_min_max_all_nulls() {
        let a = StringArray::from(vec![None, None]);
        assert_eq!(None, min_string(&a));
        assert_eq!(None, max_string(&a));
    }

    #[test]
    fn test_string_min_max_1() {
        let a = StringArray::from(vec![None, None, Some("b"), Some("a")]);
        assert_eq!(Some("a"), min_string(&a));
        assert_eq!(Some("b"), max_string(&a));
    }
}
<commit_msg>ARROW-10767: [Rust] Speed up sum with nulls (non-simd)<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Defines aggregations over Arrow arrays.

use std::ops::Add;

use crate::array::{Array, GenericStringArray, PrimitiveArray, StringOffsetSizeTrait};
use crate::datatypes::ArrowNumericType;

/// Helper macro to perform min/max of strings
fn min_max_string<T: StringOffsetSizeTrait, F: Fn(&str, &str) -> bool>(
    array: &GenericStringArray<T>,
    cmp: F,
) -> Option<&str> {
    let null_count = array.null_count();

    if null_count == array.len() {
        return None;
    }
    let data = array.data();
    let mut n;
    if null_count == 0 {
        n = array.value(0);
        for i in 1..data.len() {
            let item = array.value(i);
            if cmp(&n, item) {
                n = item;
            }
        }
    } else {
        n = "";
        let mut has_value = false;

        for i in 0..data.len() {
            let item = array.value(i);
            if data.is_valid(i) && (!has_value || cmp(&n, item)) {
                has_value = true;
                n = item;
            }
        }
    }
    Some(n)
}

/// Returns the minimum value in the array, according to the natural order.
pub fn min<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
where
    T: ArrowNumericType,
{
    min_max_helper(array, |a, b| a > b)
}

/// Returns the maximum value in the array, according to the natural order.
pub fn max<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
where
    T: ArrowNumericType,
{
    min_max_helper(array, |a, b| a < b)
}

/// Returns the maximum value in the string array, according to the natural order.
pub fn max_string<T: StringOffsetSizeTrait>(
    array: &GenericStringArray<T>,
) -> Option<&str> {
    min_max_string(array, |a, b| a < b)
}

/// Returns the minimum value in the string array, according to the natural order.
pub fn min_string<T: StringOffsetSizeTrait>(
    array: &GenericStringArray<T>,
) -> Option<&str> {
    min_max_string(array, |a, b| a > b)
}

/// Helper function to perform min/max lambda function on values from a numeric array.
fn min_max_helper<T, F>(array: &PrimitiveArray<T>, cmp: F) -> Option<T::Native>
where
    T: ArrowNumericType,
    F: Fn(&T::Native, &T::Native) -> bool,
{
    let null_count = array.null_count();

    // Includes case array.len() == 0
    if null_count == array.len() {
        return None;
    }

    let data = array.data();
    let m = array.value_slice(0, data.len());
    let mut n;

    if null_count == 0 {
        // optimized path for arrays without null values
        n = m[1..]
            .iter()
            .fold(m[0], |max, item| if cmp(&max, item) { *item } else { max });
    } else {
        n = T::default_value();
        let mut has_value = false;
        for (i, item) in m.iter().enumerate() {
            if data.is_valid(i) && (!has_value || cmp(&n, item)) {
                has_value = true;
                n = *item
            }
        }
    }
    Some(n)
}

/// Returns the sum of values in the array.
///
/// Returns `None` if the array is empty or only contains null values.
#[cfg(not(all(any(target_arch = "x86", target_arch = "x86_64"), feature = "simd")))]
pub fn sum<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
where
    T: ArrowNumericType,
    T::Native: Add<Output = T::Native>,
{
    let null_count = array.null_count();

    if null_count == array.len() {
        return None;
    }

    let data: &[T::Native] = array.value_slice(0, array.len());

    match array.data().null_buffer() {
        None => {
            let sum = data.iter().fold(T::default_value(), |accumulator, value| {
                accumulator + *value
            });

            Some(sum)
        }
        Some(buffer) => {
            let mut sum = T::default_value();
            let data_chunks = data.chunks_exact(64);
            let remainder = data_chunks.remainder();

            let bit_chunks = buffer.bit_chunks(array.offset(), array.len());
            data_chunks
                .zip(bit_chunks.iter())
                .for_each(|(chunk, mask)| {
                    // index_mask has value 1 << i in the loop
                    let mut index_mask = 1;
                    chunk.iter().for_each(|value| {
                        if (mask & index_mask) != 0 {
                            sum = sum + *value;
                        }
                        index_mask <<= 1;
                    });
                });

            let remainder_bits = bit_chunks.remainder_bits();

            remainder.iter().enumerate().for_each(|(i, value)| {
                if remainder_bits & (1 << i) != 0 {
                    sum = sum + *value;
                }
            });

            Some(sum)
        }
    }
}

/// Returns the sum of values in the array.
///
/// Returns `None` if the array is empty or only contains null values.
#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), feature = "simd"))]
pub fn sum<T: ArrowNumericType>(array: &PrimitiveArray<T>) -> Option<T::Native>
where
    T::Native: Add<Output = T::Native>,
{
    let null_count = array.null_count();

    if null_count == array.len() {
        return None;
    }

    let data: &[T::Native] = array.value_slice(0, array.len());

    let mut vector_sum = T::init(T::default_value());
    let mut remainder_sum = T::default_value();

    match array.data().null_buffer() {
        None => {
            let data_chunks = data.chunks_exact(64);
            let remainder = data_chunks.remainder();

            data_chunks.for_each(|chunk| {
                chunk.chunks_exact(T::lanes()).for_each(|chunk| {
                    let chunk = T::load(&chunk);
                    vector_sum = vector_sum + chunk;
                });
            });

            remainder.iter().for_each(|value| {
                remainder_sum = remainder_sum + *value;
            });
        }
        Some(buffer) => {
            // process data in chunks of 64 elements since we also get 64 bits of validity information at a time
            let data_chunks = data.chunks_exact(64);
            let remainder = data_chunks.remainder();

            let bit_chunks = buffer.bit_chunks(array.offset(), array.len());
            let remainder_bits = bit_chunks.remainder_bits();

            data_chunks.zip(bit_chunks).for_each(|(chunk, mut mask)| {
                // split chunks further into slices corresponding to the vector length
                // the compiler is able to unroll this inner loop and remove bounds checks
                // since the outer chunk size (64) is always a multiple of the number of lanes
                chunk.chunks_exact(T::lanes()).for_each(|chunk| {
                    let zero = T::init(T::default_value());
                    let vecmask = T::mask_from_u64(mask);
                    let chunk = T::load(&chunk);
                    let blended = T::mask_select(vecmask, chunk, zero);

                    vector_sum = vector_sum + blended;

                    mask = mask >> T::lanes();
                });
            });

            remainder.iter().enumerate().for_each(|(i, value)| {
                if remainder_bits & (1 << i) != 0 {
                    remainder_sum = remainder_sum + *value;
                }
            });
        }
    }

    // calculate horizontal sum of accumulator by writing to a temporary
    // this is probably faster than extracting individual lanes
    // the compiler is free to optimize this to something faster
    let tmp = &mut [T::default_value(); 64];
    T::write(vector_sum, &mut tmp[0..T::lanes()]);

    let mut total_sum = T::default_value();
    tmp[0..T::lanes()]
        .iter()
        .for_each(|lane| total_sum = total_sum + *lane);

    total_sum = total_sum + remainder_sum;

    Some(total_sum)
}

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

    #[test]
    fn test_primitive_array_sum() {
        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
        assert_eq!(15, sum(&a).unwrap());
    }

    #[test]
    fn test_primitive_array_float_sum() {
        let a = Float64Array::from(vec![1.1, 2.2, 3.3, 4.4, 5.5]);
        assert!(16.5 - sum(&a).unwrap() < f64::EPSILON);
    }

    #[test]
    fn test_primitive_array_sum_with_nulls() {
        let a = Int32Array::from(vec![None, Some(2), Some(3), None, Some(5)]);
        assert_eq!(10, sum(&a).unwrap());
    }

    #[test]
    fn test_primitive_array_sum_all_nulls() {
        let a = Int32Array::from(vec![None, None, None]);
        assert_eq!(None, sum(&a));
    }

    #[test]
    fn test_buffer_array_min_max() {
        let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
        assert_eq!(5, min(&a).unwrap());
        assert_eq!(9, max(&a).unwrap());
    }

    #[test]
    fn test_buffer_array_min_max_with_nulls() {
        let a = Int32Array::from(vec![Some(5), None, None, Some(8), Some(9)]);
        assert_eq!(5, min(&a).unwrap());
        assert_eq!(9, max(&a).unwrap());
    }

    #[test]
    fn test_buffer_min_max_1() {
        let a = Int32Array::from(vec![None, None, Some(5), Some(2)]);
        assert_eq!(Some(2), min(&a));
        assert_eq!(Some(5), max(&a));
    }

    #[test]
    fn test_string_min_max_with_nulls() {
        let a = StringArray::from(vec![Some("b"), None, None, Some("a"), Some("c")]);
        assert_eq!("a", min_string(&a).unwrap());
        assert_eq!("c", max_string(&a).unwrap());
    }

    #[test]
    fn test_string_min_max_all_nulls() {
        let a = StringArray::from(vec![None, None]);
        assert_eq!(None, min_string(&a));
        assert_eq!(None, max_string(&a));
    }

    #[test]
    fn test_string_min_max_1() {
        let a = StringArray::from(vec![None, None, Some("b"), Some("a")]);
        assert_eq!(Some("a"), min_string(&a));
        assert_eq!(Some("b"), max_string(&a));
    }
}
<|endoftext|>