<commit_before>// Higher level OpenCL wrappers.

use CL::*;
use CL::ll::*;
use error::check;

struct Platform {
    id: cl_platform_id
}

impl Platform {
    fn get_devices() -> ~[Device] {
        get_devices(self)
    }

    fn name() -> ~str {
        let mut size = 0;

        clGetPlatformInfo(self.id,
                          CL_PLATFORM_NAME,
                          0,
                          ptr::null(),
                          ptr::mut_addr_of(&size));

        let name = str::repeat(" ", size as uint);

        do str::as_buf(name) |p, len| {
            clGetPlatformInfo(self.id,
                              CL_PLATFORM_NAME,
                              len as libc::size_t,
                              p as *libc::c_void,
                              ptr::mut_addr_of(&size));
        };

        move name
    }
}

pub fn get_platforms() -> ~[Platform] {
    let mut num_platforms = 0;

    // TODO: Check result status
    clGetPlatformIDs(0, ptr::null(), ptr::addr_of(&num_platforms));

    let ids = vec::to_mut(vec::from_elem(num_platforms as uint,
                                         0 as cl_platform_id));

    do vec::as_imm_buf(ids) |ids, len| {
        clGetPlatformIDs(len as cl_uint,
                         ids, ptr::addr_of(&num_platforms))
    };

    do ids.map |id| { Platform { id: *id } }
}

struct Device {
    id: cl_device_id
}

pub fn get_devices(platform: Platform) -> ~[Device] {
    let mut num_devices = 0;

    clGetDeviceIDs(platform.id, CL_DEVICE_TYPE_GPU, 0, ptr::null(),
                   ptr::addr_of(&num_devices));

    let ids = vec::to_mut(vec::from_elem(num_devices as uint,
                                         0 as cl_device_id));
    do vec::as_imm_buf(ids) |ids, len| {
        clGetDeviceIDs(platform.id, CL_DEVICE_TYPE_GPU, len as cl_uint,
                       ids, ptr::addr_of(&num_devices));
    };

    do ids.map |id| { Device {id: *id }}
}

struct Context {
    ctx: cl_context,

    drop {
        clReleaseContext(self.ctx);
    }
}

pub fn create_context(device: Device) -> Context {
    // TODO: Support for multiple devices
    let mut errcode = 0;

    // TODO: Proper error messages
    let ctx = clCreateContext(ptr::null(), 1, ptr::addr_of(&device.id),
                              ptr::null(), ptr::null(), ptr::addr_of(&errcode));

    check(errcode, "Failed to create opencl context!");

    Context { ctx: ctx }
}

struct CommandQueue {
    cqueue: cl_command_queue,
    device: Device,

    drop {
        clReleaseCommandQueue(self.cqueue);
    }
}

pub fn create_commandqueue(ctx: & Context, device: Device) -> CommandQueue {
    let mut errcode = 0;

    let cqueue = clCreateCommandQueue(ctx.ctx, device.id, 0, ptr::addr_of(&errcode));

    check(errcode, "Failed to create command queue!");

    CommandQueue {
        cqueue: cqueue,
        device: device
    }
}

struct Buffer {
    buffer: cl_mem,
    size: int,

    drop {
        clReleaseMemObject(self.buffer);
    }
}

// TODO: How to make this function cleaner and nice
pub fn create_buffer(ctx: & Context, size: int, flags: cl_mem_flags) -> Buffer {
    let mut errcode = 0;

    let buffer = clCreateBuffer(ctx.ctx, flags, size as libc::size_t, ptr::null(),
                                ptr::addr_of(&errcode));

    check(errcode, "Failed to create buffer!");

    Buffer { buffer: buffer, size: size }
}

pub fn enqueue_write_buffer<T: KernelArg>(
    cqueue: & CommandQueue,
    buf: & Buffer,
    host_vector: T) unsafe
{
    let ret = clEnqueueWriteBuffer(
        cqueue.cqueue, buf.buffer, CL_TRUE, 0,
        buf.size as libc::size_t,
        host_vector.get_value(),
        0, ptr::null(), ptr::null());
    check(ret, "Failed to enqueue write buffer!");
}

pub fn enqueue_read_buffer<T: KernelArg>(
    cqueue: & CommandQueue,
    buf: & Buffer,
    host_vector: T) unsafe
{
    let ret = clEnqueueReadBuffer(
        cqueue.cqueue, buf.buffer, CL_TRUE, 0,
        buf.size as libc::size_t,
        host_vector.get_value(), 0, ptr::null(), ptr::null());

    check(ret, "Failed to enqueue read buffer!");
}

struct Program {
    prg: cl_program,

    drop {
        clReleaseProgram(self.prg);
    }
}

pub impl Program {
    fn build(&self, device: Device) {
        build_program(self, device);
    }
}

// TODO: Support multiple devices
pub fn create_program_with_binary(ctx: & Context, device: Device,
                                  binary_path: & Path) -> Program {
    let mut errcode = 0;
    let binary = match move io::read_whole_file_str(binary_path) {
        result::Ok(move binary) => move binary,
        Err(e) => fail fmt!("%?", e)
    };
    let program = do str::as_c_str(binary) |kernel_binary| {
        clCreateProgramWithBinary(ctx.ctx, 1, ptr::addr_of(&device.id),
                                  ptr::addr_of(&(binary.len() + 1)) as *libc::size_t,
                                  ptr::addr_of(&kernel_binary) as **libc::c_uchar,
                                  ptr::null(),
                                  ptr::addr_of(&errcode))
    };

    check(errcode, "Failed to create open cl program with binary!");

    Program { prg: program }
}

pub fn build_program(program: & Program, device: Device) {
    let ret = clBuildProgram(program.prg, 1, ptr::addr_of(&device.id),
                             ptr::null(), ptr::null(), ptr::null());
    if ret != CL_SUCCESS as cl_int {
        let mut logv = ~"";
        for uint::range(0,1024*1204) |_i| {
            str::push_char(& mut logv, ' ');
        }
        do str::as_buf(logv) |logs, l| {
            let r = clGetProgramBuildInfo(program.prg, device.id, CL_PROGRAM_BUILD_LOG, l as libc::size_t, logs as *libc::c_void, ptr::null());
            check(r, "failed to get build info!");
        }
        info!("%s", logv);
        //io::println(logv);
        fail fmt!("Failure during program building! %?", ret)
    }
}

struct Kernel {
    kernel: cl_kernel,

    drop {
        clReleaseKernel(self.kernel);
    }
}

impl Kernel {
    fn set_arg<T: KernelArg>(i: uint, x: &T) {
        set_kernel_arg(&self, i as CL::cl_uint, x)
    }
}

pub fn create_kernel(program: & Program, kernel: & str) -> Kernel unsafe{
    let mut errcode = 0;
    let bytes = str::to_bytes(kernel);
    let kernel = clCreateKernel(program.prg, vec::raw::to_ptr(bytes) as *libc::c_char, ptr::addr_of(&errcode));

    check(errcode, "Failed to create kernel!");

    Kernel { kernel: kernel }
}

pub trait KernelArg{
    fn get_value() -> *libc::c_void;
}

pub fn set_kernel_arg<T: KernelArg>(kernel: & Kernel, position: cl_uint, arg: & T) unsafe{
    // TODO: How to set different argument types. Currently only support cl_mem
    let ret = clSetKernelArg(kernel.kernel, position,
                             sys::size_of::<cl_mem>() as libc::size_t,
                             arg.get_value());

    check(ret, "Failed to set kernel arg!");
}

pub fn enqueue_nd_range_kernel(cqueue: & CommandQueue, kernel: & Kernel, work_dim: cl_uint,
                               _global_work_offset: int, global_work_size: int,
                               local_work_size: int) unsafe{
    let ret = clEnqueueNDRangeKernel(cqueue.cqueue, kernel.kernel, work_dim,
                                     // ptr::addr_of(&global_work_offset) as *libc::size_t,
                                     ptr::null(),
                                     ptr::addr_of(&global_work_size) as *libc::size_t,
                                     ptr::addr_of(&local_work_size) as *libc::size_t,
                                     0, ptr::null(), ptr::null());
    check(ret, "Failed to enqueue nd range kernel!");
}

pub impl<T> &~[T]: KernelArg {
    fn get_value() -> *libc::c_void {
        do vec::as_imm_buf(*self) |p, _len| {
            p as *libc::c_void
        }
    }
}

pub impl<T> &~[mut T]: KernelArg {
    fn get_value() -> *libc::c_void {
        do vec::as_imm_buf(*self) |p, _len| {
            p as *libc::c_void
        }
    }
}

pub impl Buffer: KernelArg{
    fn get_value() -> *libc::c_void{
        ptr::addr_of(&self.buffer) as *libc::c_void
    }
}

pub impl *libc::c_void: KernelArg{
    fn get_value() -> *libc::c_void {
        self
    }
}

/**
This packages an OpenCL context, a device, and a command queue to
simplify handling all these structures.
*/
struct ComputeContext {
    ctx: Context,
    device: Device,
    q: CommandQueue
}

impl ComputeContext {
    fn create_program_from_source(src: &str) -> Program {
        do str::as_c_str(src) |src| {
            let mut status = CL_SUCCESS as cl_int;
            let program = clCreateProgramWithSource(
                self.ctx.ctx,
                1,
                ptr::addr_of(&src),
                ptr::null(),
                ptr::addr_of(&status));
            check(status, "Could not create program");

            Program { prg: program }
        }
    }
}

pub fn create_compute_context() -> @ComputeContext {
    // Enumerate all platforms until we find a device that works.

    for get_platforms().each |p| {
        let devices = p.get_devices();
        if devices.len() > 0 {
            let device = devices[0];
            let ctx = create_context(device);
            let q = create_commandqueue(&ctx, device);
            return @ComputeContext {
                ctx: move ctx,
                device: move device,
                q: move q
            }
        }
    }
    fail ~"Could not find an acceptable device."
}

#[cfg(test)]
mod test {
    #[test]
    fn program_build() {
        let src = "__kernel void test(__global int *i) { \
                       *i += 1; \
                   }";
        let ctx = create_compute_context();
        let prog = ctx.create_program_from_source(src);
        prog.build(ctx.device);
    }
}<commit_msg>Better reporting of build errors.<commit_after>// Higher level OpenCL wrappers.

use CL::*;
use CL::ll::*;
use error::check;

struct Platform {
    id: cl_platform_id
}

impl Platform {
    fn get_devices() -> ~[Device] {
        get_devices(self)
    }

    fn name() -> ~str {
        let mut size = 0;

        clGetPlatformInfo(self.id,
                          CL_PLATFORM_NAME,
                          0,
                          ptr::null(),
                          ptr::mut_addr_of(&size));

        let name = str::repeat(" ", size as uint);

        do str::as_buf(name) |p, len| {
            clGetPlatformInfo(self.id,
                              CL_PLATFORM_NAME,
                              len as libc::size_t,
                              p as *libc::c_void,
                              ptr::mut_addr_of(&size));
        };

        move name
    }
}

pub fn get_platforms() -> ~[Platform] {
    let mut num_platforms = 0;

    // TODO: Check result status
    clGetPlatformIDs(0, ptr::null(), ptr::addr_of(&num_platforms));

    let ids = vec::to_mut(vec::from_elem(num_platforms as uint,
                                         0 as cl_platform_id));

    do vec::as_imm_buf(ids) |ids, len| {
        clGetPlatformIDs(len as cl_uint,
                         ids, ptr::addr_of(&num_platforms))
    };

    do ids.map |id| { Platform { id: *id } }
}

struct Device {
    id: cl_device_id
}

pub fn get_devices(platform: Platform) -> ~[Device] {
    let mut num_devices = 0;

    clGetDeviceIDs(platform.id, CL_DEVICE_TYPE_GPU, 0, ptr::null(),
                   ptr::addr_of(&num_devices));

    let ids = vec::to_mut(vec::from_elem(num_devices as uint,
                                         0 as cl_device_id));
    do vec::as_imm_buf(ids) |ids, len| {
        clGetDeviceIDs(platform.id, CL_DEVICE_TYPE_GPU, len as cl_uint,
                       ids, ptr::addr_of(&num_devices));
    };

    do ids.map |id| { Device {id: *id }}
}

struct Context {
    ctx: cl_context,

    drop {
        clReleaseContext(self.ctx);
    }
}

pub fn create_context(device: Device) -> Context {
    // TODO: Support for multiple devices
    let mut errcode = 0;

    // TODO: Proper error messages
    let ctx = clCreateContext(ptr::null(), 1, ptr::addr_of(&device.id),
                              ptr::null(), ptr::null(), ptr::addr_of(&errcode));

    check(errcode, "Failed to create opencl context!");

    Context { ctx: ctx }
}

struct CommandQueue {
    cqueue: cl_command_queue,
    device: Device,

    drop {
        clReleaseCommandQueue(self.cqueue);
    }
}

pub fn create_commandqueue(ctx: & Context, device: Device) -> CommandQueue {
    let mut errcode = 0;

    let cqueue = clCreateCommandQueue(ctx.ctx, device.id, 0, ptr::addr_of(&errcode));

    check(errcode, "Failed to create command queue!");

    CommandQueue {
        cqueue: cqueue,
        device: device
    }
}

struct Buffer {
    buffer: cl_mem,
    size: int,

    drop {
        clReleaseMemObject(self.buffer);
    }
}

// TODO: How to make this function cleaner and nice
pub fn create_buffer(ctx: & Context, size: int, flags: cl_mem_flags) -> Buffer {
    let mut errcode = 0;

    let buffer = clCreateBuffer(ctx.ctx, flags, size as libc::size_t, ptr::null(),
                                ptr::addr_of(&errcode));

    check(errcode, "Failed to create buffer!");

    Buffer { buffer: buffer, size: size }
}

pub fn enqueue_write_buffer<T: KernelArg>(
    cqueue: & CommandQueue,
    buf: & Buffer,
    host_vector: T) unsafe
{
    let ret = clEnqueueWriteBuffer(
        cqueue.cqueue, buf.buffer, CL_TRUE, 0,
        buf.size as libc::size_t,
        host_vector.get_value(),
        0, ptr::null(), ptr::null());
    check(ret, "Failed to enqueue write buffer!");
}

pub fn enqueue_read_buffer<T: KernelArg>(
    cqueue: & CommandQueue,
    buf: & Buffer,
    host_vector: T) unsafe
{
    let ret = clEnqueueReadBuffer(
        cqueue.cqueue, buf.buffer, CL_TRUE, 0,
        buf.size as libc::size_t,
        host_vector.get_value(), 0, ptr::null(), ptr::null());

    check(ret, "Failed to enqueue read buffer!");
}

struct Program {
    prg: cl_program,

    drop {
        clReleaseProgram(self.prg);
    }
}

pub impl Program {
    fn build(&self, device: Device) -> Result<(), ~str> {
        build_program(self, device)
    }
}

// TODO: Support multiple devices
pub fn create_program_with_binary(ctx: & Context, device: Device,
                                  binary_path: & Path) -> Program {
    let mut errcode = 0;
    let binary = match move io::read_whole_file_str(binary_path) {
        result::Ok(move binary) => move binary,
        Err(e) => fail fmt!("%?", e)
    };
    let program = do str::as_c_str(binary) |kernel_binary| {
        clCreateProgramWithBinary(ctx.ctx, 1, ptr::addr_of(&device.id),
                                  ptr::addr_of(&(binary.len() + 1)) as *libc::size_t,
                                  ptr::addr_of(&kernel_binary) as **libc::c_uchar,
                                  ptr::null(),
                                  ptr::addr_of(&errcode))
    };

    check(errcode, "Failed to create open cl program with binary!");

    Program { prg: program }
}

pub fn build_program(program: & Program, device: Device) -> Result<(), ~str> {
    let ret = clBuildProgram(program.prg, 1, ptr::addr_of(&device.id),
                             ptr::null(), ptr::null(), ptr::null());
    if ret == CL_SUCCESS as cl_int {
        Ok(())
    }
    else {
        let mut size = 0 as libc::size_t;
        let status = clGetProgramBuildInfo(
            program.prg,
            device.id,
            CL_PROGRAM_BUILD_LOG,
            0,
            ptr::null(),
            ptr::addr_of(&size));
        check(status, "Could not get build log");

        let buf = vec::from_elem(size as uint, 0u8);
        do vec::as_imm_buf(buf) |p, len| {
            let status = clGetProgramBuildInfo(
                program.prg,
                device.id,
                CL_PROGRAM_BUILD_LOG,
                len as libc::size_t,
                p as *libc::c_void,
                ptr::null());
            check(status, "Could not get build log");

            unsafe { Err(str::raw::from_c_str(p as *libc::c_char)) }
        }
    }
}

struct Kernel {
    kernel: cl_kernel,

    drop {
        clReleaseKernel(self.kernel);
    }
}

impl Kernel {
    fn set_arg<T: KernelArg>(i: uint, x: &T) {
        set_kernel_arg(&self, i as CL::cl_uint, x)
    }
}

pub fn create_kernel(program: & Program, kernel: & str) -> Kernel unsafe{
    let mut errcode = 0;
    let bytes = str::to_bytes(kernel);
    let kernel = clCreateKernel(program.prg, vec::raw::to_ptr(bytes) as *libc::c_char, ptr::addr_of(&errcode));

    check(errcode, "Failed to create kernel!");

    Kernel { kernel: kernel }
}

pub trait KernelArg{
    fn get_value() -> *libc::c_void;
}

pub fn set_kernel_arg<T: KernelArg>(kernel: & Kernel, position: cl_uint, arg: & T) unsafe{
    // TODO: How to set different argument types. Currently only support cl_mem
    let ret = clSetKernelArg(kernel.kernel, position,
                             sys::size_of::<cl_mem>() as libc::size_t,
                             arg.get_value());

    check(ret, "Failed to set kernel arg!");
}

pub fn enqueue_nd_range_kernel(cqueue: & CommandQueue, kernel: & Kernel, work_dim: cl_uint,
                               _global_work_offset: int, global_work_size: int,
                               local_work_size: int) unsafe{
    let ret = clEnqueueNDRangeKernel(cqueue.cqueue, kernel.kernel, work_dim,
                                     // ptr::addr_of(&global_work_offset) as *libc::size_t,
                                     ptr::null(),
                                     ptr::addr_of(&global_work_size) as *libc::size_t,
                                     ptr::addr_of(&local_work_size) as *libc::size_t,
                                     0, ptr::null(), ptr::null());
    check(ret, "Failed to enqueue nd range kernel!");
}

pub impl<T> &~[T]: KernelArg {
    fn get_value() -> *libc::c_void {
        do vec::as_imm_buf(*self) |p, _len| {
            p as *libc::c_void
        }
    }
}

pub impl<T> &~[mut T]: KernelArg {
    fn get_value() -> *libc::c_void {
        do vec::as_imm_buf(*self) |p, _len| {
            p as *libc::c_void
        }
    }
}

pub impl Buffer: KernelArg{
    fn get_value() -> *libc::c_void{
        ptr::addr_of(&self.buffer) as *libc::c_void
    }
}

pub impl *libc::c_void: KernelArg{
    fn get_value() -> *libc::c_void {
        self
    }
}

/**
This packages an OpenCL context, a device, and a command queue to
simplify handling all these structures.
*/
struct ComputeContext {
    ctx: Context,
    device: Device,
    q: CommandQueue
}

impl ComputeContext {
    fn create_program_from_source(src: &str) -> Program {
        do str::as_c_str(src) |src| {
            let mut status = CL_SUCCESS as cl_int;
            let program = clCreateProgramWithSource(
                self.ctx.ctx,
                1,
                ptr::addr_of(&src),
                ptr::null(),
                ptr::addr_of(&status));
            check(status, "Could not create program");

            Program { prg: program }
        }
    }
}

pub fn create_compute_context() -> @ComputeContext {
    // Enumerate all platforms until we find a device that works.

    for get_platforms().each |p| {
        let devices = p.get_devices();
        if devices.len() > 0 {
            let device = devices[0];
            let ctx = create_context(device);
            let q = create_commandqueue(&ctx, device);
            return @ComputeContext {
                ctx: move ctx,
                device: move device,
                q: move q
            }
        }
    }
    fail ~"Could not find an acceptable device."
}

#[cfg(test)]
mod test {
    #[test]
    fn program_build() {
        let src = "__kernel void test(__global int *i) { \
                       *i += 1; \
                   }";
        let ctx = create_compute_context();
        let prog = ctx.create_program_from_source(src);
        prog.build(ctx.device);
    }
}<|endoftext|>


<commit_before>use mio::{Events, Poll, PollOpt, Ready, Registration, SetReadiness, Token};
use mio::event::Evented;
use std::time::Duration;

#[test]
fn smoke() {
    let poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(128);

    let (r, set) = Registration::new2();
    r.register(&poll, Token(0), Ready::readable(), PollOpt::edge()).unwrap();

    let n = poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();
    assert_eq!(n, 0);

    set.set_readiness(Ready::readable()).unwrap();

    let n = poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();
    assert_eq!(n, 1);

    assert_eq!(events.get(0).unwrap().token(), Token(0));
}

#[test]
fn set_readiness_before_register() {
    use std::sync::{Arc, Barrier};
    use std::thread;

    let poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(128);

    for _ in 0..5_000 {
        let (r, set) = Registration::new2();

        let b1 = Arc::new(Barrier::new(2));
        let b2 = b1.clone();

        let th = thread::spawn(move || {
            // set readiness before register
            set.set_readiness(Ready::readable()).unwrap();

            // run into barrier so both can pass
            b2.wait();
        });

        // wait for readiness
        b1.wait();

        // now register
        poll.register(&r, Token(123), Ready::readable(), PollOpt::edge()).unwrap();

        loop {
            let n = poll.poll(&mut events, None).unwrap();

            if n == 0 {
                continue;
            }

            assert_eq!(n, 1);
            assert_eq!(events.get(0).unwrap().token(), Token(123));
            break;
        }

        th.join().unwrap();
    }
}

#[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
mod stress {
    use mio::{Events, Poll, PollOpt, Ready, Registration, SetReadiness, Token};
    use mio::event::Evented;
    use std::time::Duration;

    #[test]
    fn single_threaded_poll() {
        use std::sync::Arc;
        use std::sync::atomic::AtomicUsize;
        use std::sync::atomic::Ordering::{Acquire, Release};
        use std::thread;

        const NUM_ATTEMPTS: usize = 30;
        const NUM_ITERS: usize = 500;
        const NUM_THREADS: usize = 4;
        const NUM_REGISTRATIONS: usize = 128;

        for _ in 0..NUM_ATTEMPTS {
            let poll = Poll::new().unwrap();
            let mut events = Events::with_capacity(128);

            let registrations: Vec<_> = (0..NUM_REGISTRATIONS).map(|i| {
                let (r, s) = Registration::new2();
                r.register(&poll, Token(i), Ready::readable(), PollOpt::edge()).unwrap();
                (r, s)
            }).collect();

            let mut ready: Vec<_> = (0..NUM_REGISTRATIONS).map(|_| Ready::empty()).collect();

            let remaining = Arc::new(AtomicUsize::new(NUM_THREADS));

            for _ in 0..NUM_THREADS {
                let remaining = remaining.clone();

                let set_readiness: Vec<SetReadiness> =
                    registrations.iter().map(|r| r.1.clone()).collect();

                thread::spawn(move || {
                    for _ in 0..NUM_ITERS {
                        for i in 0..NUM_REGISTRATIONS {
                            set_readiness[i].set_readiness(Ready::readable()).unwrap();
                            set_readiness[i].set_readiness(Ready::empty()).unwrap();
                            set_readiness[i].set_readiness(Ready::writable()).unwrap();
                            set_readiness[i].set_readiness(Ready::readable() | Ready::writable()).unwrap();
                            set_readiness[i].set_readiness(Ready::empty()).unwrap();
                        }
                    }

                    for i in 0..NUM_REGISTRATIONS {
                        set_readiness[i].set_readiness(Ready::readable()).unwrap();
                    }

                    remaining.fetch_sub(1, Release);
                });
            }

            while remaining.load(Acquire) > 0 {
                // Set interest
                for (i, &(ref r, _)) in registrations.iter().enumerate() {
                    r.reregister(&poll, Token(i), Ready::writable(), PollOpt::edge()).unwrap();
                }

                poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();

                for event in &events {
                    ready[event.token().0] = event.readiness();
                }

                // Update registration
                // Set interest
                for (i, &(ref r, _)) in registrations.iter().enumerate() {
                    r.reregister(&poll, Token(i), Ready::readable(), PollOpt::edge()).unwrap();
                }
            }

            // Finall polls, repeat until readiness-queue empty
            loop {
                // Might not read all events from custom-event-queue at once, implementation dependend
                poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();
                if events.is_empty() {
                    // no more events in readiness queue pending
                    break;
                }
                for event in &events {
                    ready[event.token().0] = event.readiness();
                }
            }

            // Everything should be flagged as readable
            for ready in ready {
                assert_eq!(ready, Ready::readable());
            }
        }
    }

    #[test]
    fn multi_threaded_poll() {
        use std::sync::{Arc, Barrier};
        use std::sync::atomic::{AtomicUsize};
        use std::sync::atomic::Ordering::{Relaxed, SeqCst};
        use std::thread;

        const ENTRIES: usize = 10_000;
        const PER_ENTRY: usize = 16;
        const THREADS: usize = 4;
        const NUM: usize = ENTRIES * PER_ENTRY;

        struct Entry {
            #[allow(dead_code)]
            registration: Registration,
            set_readiness: SetReadiness,
            num: AtomicUsize,
        }

        impl Entry {
            fn fire(&self) {
                self.set_readiness.set_readiness(Ready::readable()).unwrap();
            }
        }

        let poll = Arc::new(Poll::new().unwrap());
        let mut entries = vec![];

        // Create entries
        for i in 0..ENTRIES {
            let (registration, set_readiness) = Registration::new2();
            registration.register(&poll, Token(i), Ready::readable(), PollOpt::edge()).unwrap();

            entries.push(Entry {
                registration: registration,
                set_readiness: set_readiness,
                num: AtomicUsize::new(0),
            });
        }

        let total = Arc::new(AtomicUsize::new(0));
        let entries = Arc::new(entries);
        let barrier = Arc::new(Barrier::new(THREADS));

        let mut threads = vec![];

        for th in 0..THREADS {
            let poll = poll.clone();
            let total = total.clone();
            let entries = entries.clone();
            let barrier = barrier.clone();

            threads.push(thread::spawn(move || {
                let mut events = Events::with_capacity(128);

                barrier.wait();

                // Prime all the registrations
                let mut i = th;
                while i < ENTRIES {
                    entries[i].fire();
                    i += THREADS;
                }

                let mut n = 0;

                while total.load(SeqCst) < NUM {
                    // A poll timeout is necessary here because there may be more
                    // than one threads blocked in `poll` when the final wakeup
                    // notification arrives (and only notifies one thread).
                    n += poll.poll(&mut events, Some(Duration::from_millis(100))).unwrap();

                    let mut num_this_tick = 0;

                    for event in &events {
                        let e = &entries[event.token().0];

                        let mut num = e.num.load(Relaxed);

                        loop {
                            if num < PER_ENTRY {
                                let actual = e.num.compare_and_swap(num, num + 1, Relaxed);

                                if actual == num {
                                    num_this_tick += 1;
                                    e.fire();
                                    break;
                                }

                                num = actual;
                            } else {
                                break;
                            }
                        }
                    }

                    total.fetch_add(num_this_tick, SeqCst);
                }

                n
            }));
        }

        let per_thread: Vec<_> = threads.into_iter()
            .map(|th| th.join().unwrap())
            .collect();

        for entry in entries.iter() {
            assert_eq!(PER_ENTRY, entry.num.load(Relaxed));
        }

        for th in per_thread {
            // Kind of annoying that we can't really test anything better than this,
            // but CI tends to be very non deterministic when it comes to multi
            // threading.
            assert!(th > 0, "actual={:?}", th);
        }
    }

    #[test]
    fn with_small_events_collection() {
        const N: usize = 8;
        const ITER: usize = 1_000;

        use std::sync::{Arc, Barrier};
        use std::sync::atomic::AtomicBool;
        use std::sync::atomic::Ordering::{Acquire, Release};
        use std::thread;

        let poll = Poll::new().unwrap();
        let mut registrations = vec![];

        let barrier = Arc::new(Barrier::new(N + 1));
        let done = Arc::new(AtomicBool::new(false));

        for i in 0..N {
            let (registration, set_readiness) = Registration::new2();
            poll.register(&registration, Token(i), Ready::readable(), PollOpt::edge()).unwrap();

            registrations.push(registration);

            let barrier = barrier.clone();
            let done = done.clone();

            thread::spawn(move || {
                barrier.wait();

                while !done.load(Acquire) {
                    set_readiness.set_readiness(Ready::readable()).unwrap();
                }

                // Set one last time
                set_readiness.set_readiness(Ready::readable()).unwrap();
            });
        }

        let mut events = Events::with_capacity(4);

        barrier.wait();

        for _ in 0..ITER {
            poll.poll(&mut events, None).unwrap();
        }

        done.store(true, Release);

        let mut final_ready = vec![false; N];

        for _ in 0..5 {
            poll.poll(&mut events, None).unwrap();

            for event in &events {
                final_ready[event.token().0] = true;
            }

            if final_ready.iter().all(|v| *v) {
                return;
            }

            thread::sleep(Duration::from_millis(10));
        }

        panic!("dead lock?");
    }
}

#[test]
fn drop_registration_from_non_main_thread() {
    use std::thread;
    use std::sync::mpsc::channel;

    const THREADS: usize = 8;
    const ITERS: usize = 50_000;

    let mut poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(1024);
    let mut senders = Vec::with_capacity(THREADS);
    let mut token_index = 0;

    // spawn threads, which will send messages to single receiver
    for _ in 0..THREADS {
        let (tx, rx) = channel::<(Registration, SetReadiness)>();
        senders.push(tx);

        thread::spawn(move || {
            for (registration, set_readiness) in rx {
                let _ = set_readiness.set_readiness(Ready::readable());
                drop(registration);
                drop(set_readiness);
            }
        });
    }

    let mut index: usize = 0;
    for _ in 0..ITERS {
        let (registration, set_readiness) = Registration::new2();
        registration.register(&mut poll, Token(token_index), Ready::readable(), PollOpt::edge()).unwrap();
        let _ = senders[index].send((registration, set_readiness));

        token_index += 1;
        index += 1;
        if index == THREADS {
            index = 0;

            let (registration, set_readiness) = Registration::new2();
            registration.register(&mut poll, Token(token_index), Ready::readable(), PollOpt::edge()).unwrap();
            let _ = set_readiness.set_readiness(Ready::readable());
            drop(registration);
            drop(set_readiness);
            token_index += 1;

            thread::park_timeout(Duration::from_millis(0));
            let _ = poll.poll(&mut events, None).unwrap();
        }
    }
}
<commit_msg>refactor, use constant instead of magic number (#787)<commit_after>use mio::{Events, Poll, PollOpt, Ready, Registration, SetReadiness, Token};
use mio::event::Evented;
use std::time::Duration;

#[test]
fn smoke() {
    let poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(128);

    let (r, set) = Registration::new2();
    r.register(&poll, Token(0), Ready::readable(), PollOpt::edge()).unwrap();

    let n = poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();
    assert_eq!(n, 0);

    set.set_readiness(Ready::readable()).unwrap();

    let n = poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();
    assert_eq!(n, 1);

    assert_eq!(events.get(0).unwrap().token(), Token(0));
}

#[test]
fn set_readiness_before_register() {
    use std::sync::{Arc, Barrier};
    use std::thread;

    let poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(128);

    for _ in 0..5_000 {
        let (r, set) = Registration::new2();

        let b1 = Arc::new(Barrier::new(2));
        let b2 = b1.clone();

        let th = thread::spawn(move || {
            // set readiness before register
            set.set_readiness(Ready::readable()).unwrap();

            // run into barrier so both can pass
            b2.wait();
        });

        // wait for readiness
        b1.wait();

        // now register
        poll.register(&r, Token(123), Ready::readable(), PollOpt::edge()).unwrap();

        loop {
            let n = poll.poll(&mut events, None).unwrap();

            if n == 0 {
                continue;
            }

            assert_eq!(n, 1);
            assert_eq!(events.get(0).unwrap().token(), Token(123));
            break;
        }

        th.join().unwrap();
    }
}

#[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
mod stress {
    use mio::{Events, Poll, PollOpt, Ready, Registration, SetReadiness, Token};
    use mio::event::Evented;
    use std::time::Duration;

    #[test]
    fn single_threaded_poll() {
        use std::sync::Arc;
        use std::sync::atomic::AtomicUsize;
        use std::sync::atomic::Ordering::{Acquire, Release};
        use std::thread;

        const NUM_ATTEMPTS: usize = 30;
        const NUM_ITERS: usize = 500;
        const NUM_THREADS: usize = 4;
        const NUM_REGISTRATIONS: usize = 128;

        for _ in 0..NUM_ATTEMPTS {
            let poll = Poll::new().unwrap();
            let mut events = Events::with_capacity(NUM_REGISTRATIONS);

            let registrations: Vec<_> = (0..NUM_REGISTRATIONS).map(|i| {
                let (r, s) = Registration::new2();
                r.register(&poll, Token(i), Ready::readable(), PollOpt::edge()).unwrap();
                (r, s)
            }).collect();

            let mut ready: Vec<_> = (0..NUM_REGISTRATIONS).map(|_| Ready::empty()).collect();

            let remaining = Arc::new(AtomicUsize::new(NUM_THREADS));

            for _ in 0..NUM_THREADS {
                let remaining = remaining.clone();

                let set_readiness: Vec<SetReadiness> =
                    registrations.iter().map(|r| r.1.clone()).collect();

                thread::spawn(move || {
                    for _ in 0..NUM_ITERS {
                        for i in 0..NUM_REGISTRATIONS {
                            set_readiness[i].set_readiness(Ready::readable()).unwrap();
                            set_readiness[i].set_readiness(Ready::empty()).unwrap();
                            set_readiness[i].set_readiness(Ready::writable()).unwrap();
                            set_readiness[i].set_readiness(Ready::readable() | Ready::writable()).unwrap();
                            set_readiness[i].set_readiness(Ready::empty()).unwrap();
                        }
                    }

                    for i in 0..NUM_REGISTRATIONS {
                        set_readiness[i].set_readiness(Ready::readable()).unwrap();
                    }

                    remaining.fetch_sub(1, Release);
                });
            }

            while remaining.load(Acquire) > 0 {
                // Set interest
                for (i, &(ref r, _)) in registrations.iter().enumerate() {
                    r.reregister(&poll, Token(i), Ready::writable(), PollOpt::edge()).unwrap();
                }

                poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();

                for event in &events {
                    ready[event.token().0] = event.readiness();
                }

                // Update registration
                // Set interest
                for (i, &(ref r, _)) in registrations.iter().enumerate() {
                    r.reregister(&poll, Token(i), Ready::readable(), PollOpt::edge()).unwrap();
                }
            }

            // Finall polls, repeat until readiness-queue empty
            loop {
                // Might not read all events from custom-event-queue at once, implementation dependend
                poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();
                if events.is_empty() {
                    // no more events in readiness queue pending
                    break;
                }
                for event in &events {
                    ready[event.token().0] = event.readiness();
                }
            }

            // Everything should be flagged as readable
            for ready in ready {
                assert_eq!(ready, Ready::readable());
            }
        }
    }

    #[test]
    fn multi_threaded_poll() {
        use std::sync::{Arc, Barrier};
        use std::sync::atomic::{AtomicUsize};
        use std::sync::atomic::Ordering::{Relaxed, SeqCst};
        use std::thread;

        const ENTRIES: usize = 10_000;
        const PER_ENTRY: usize = 16;
        const THREADS: usize = 4;
        const NUM: usize = ENTRIES * PER_ENTRY;

        struct Entry {
            #[allow(dead_code)]
            registration: Registration,
            set_readiness: SetReadiness,
            num: AtomicUsize,
        }

        impl Entry {
            fn fire(&self) {
                self.set_readiness.set_readiness(Ready::readable()).unwrap();
            }
        }

        let poll = Arc::new(Poll::new().unwrap());
        let mut entries = vec![];

        // Create entries
        for i in 0..ENTRIES {
            let (registration, set_readiness) = Registration::new2();
            registration.register(&poll, Token(i), Ready::readable(), PollOpt::edge()).unwrap();

            entries.push(Entry {
                registration: registration,
                set_readiness: set_readiness,
                num: AtomicUsize::new(0),
            });
        }

        let total = Arc::new(AtomicUsize::new(0));
        let entries = Arc::new(entries);
        let barrier = Arc::new(Barrier::new(THREADS));

        let mut threads = vec![];

        for th in 0..THREADS {
            let poll = poll.clone();
            let total = total.clone();
            let entries = entries.clone();
            let barrier = barrier.clone();

            threads.push(thread::spawn(move || {
                let mut events = Events::with_capacity(128);

                barrier.wait();

                // Prime all the registrations
                let mut i = th;
                while i < ENTRIES {
                    entries[i].fire();
                    i += THREADS;
                }

                let mut n = 0;

                while total.load(SeqCst) < NUM {
                    // A poll timeout is necessary here because there may be more
                    // than one threads blocked in `poll` when the final wakeup
                    // notification arrives (and only notifies one thread).
                    n += poll.poll(&mut events, Some(Duration::from_millis(100))).unwrap();

                    let mut num_this_tick = 0;

                    for event in &events {
                        let e = &entries[event.token().0];

                        let mut num = e.num.load(Relaxed);

                        loop {
                            if num < PER_ENTRY {
                                let actual = e.num.compare_and_swap(num, num + 1, Relaxed);

                                if actual == num {
                                    num_this_tick += 1;
                                    e.fire();
                                    break;
                                }

                                num = actual;
                            } else {
                                break;
                            }
                        }
                    }

                    total.fetch_add(num_this_tick, SeqCst);
                }

                n
            }));
        }

        let per_thread: Vec<_> = threads.into_iter()
            .map(|th| th.join().unwrap())
            .collect();

        for entry in entries.iter() {
            assert_eq!(PER_ENTRY, entry.num.load(Relaxed));
        }

        for th in per_thread {
            // Kind of annoying that we can't really test anything better than this,
            // but CI tends to be very non deterministic when it comes to multi
            // threading.
            assert!(th > 0, "actual={:?}", th);
        }
    }

    #[test]
    fn with_small_events_collection() {
        const N: usize = 8;
        const ITER: usize = 1_000;

        use std::sync::{Arc, Barrier};
        use std::sync::atomic::AtomicBool;
        use std::sync::atomic::Ordering::{Acquire, Release};
        use std::thread;

        let poll = Poll::new().unwrap();
        let mut registrations = vec![];

        let barrier = Arc::new(Barrier::new(N + 1));
        let done = Arc::new(AtomicBool::new(false));

        for i in 0..N {
            let (registration, set_readiness) = Registration::new2();
            poll.register(&registration, Token(i), Ready::readable(), PollOpt::edge()).unwrap();

            registrations.push(registration);

            let barrier = barrier.clone();
            let done = done.clone();

            thread::spawn(move || {
                barrier.wait();

                while !done.load(Acquire) {
                    set_readiness.set_readiness(Ready::readable()).unwrap();
                }

                // Set one last time
                set_readiness.set_readiness(Ready::readable()).unwrap();
            });
        }

        let mut events = Events::with_capacity(4);

        barrier.wait();

        for _ in 0..ITER {
            poll.poll(&mut events, None).unwrap();
        }

        done.store(true, Release);

        let mut final_ready = vec![false; N];

        for _ in 0..5 {
            poll.poll(&mut events, None).unwrap();

            for event in &events {
                final_ready[event.token().0] = true;
            }

            if final_ready.iter().all(|v| *v) {
                return;
            }

            thread::sleep(Duration::from_millis(10));
        }

        panic!("dead lock?");
    }
}

#[test]
fn drop_registration_from_non_main_thread() {
    use std::thread;
    use std::sync::mpsc::channel;

    const THREADS: usize = 8;
    const ITERS: usize = 50_000;

    let mut poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(1024);
    let mut senders = Vec::with_capacity(THREADS);
    let mut token_index = 0;

    // spawn threads, which will send messages to single receiver
    for _ in 0..THREADS {
        let (tx, rx) = channel::<(Registration, SetReadiness)>();
        senders.push(tx);

        thread::spawn(move || {
            for (registration, set_readiness) in rx {
                let _ = set_readiness.set_readiness(Ready::readable());
                drop(registration);
                drop(set_readiness);
            }
        });
    }

    let mut index: usize = 0;
    for _ in 0..ITERS {
        let (registration, set_readiness) = Registration::new2();
        registration.register(&mut poll, Token(token_index), Ready::readable(), PollOpt::edge()).unwrap();
        let _ = senders[index].send((registration, set_readiness));

        token_index += 1;
        index += 1;
        if index == THREADS {
            index = 0;

            let (registration, set_readiness) = Registration::new2();
            registration.register(&mut poll, Token(token_index), Ready::readable(), PollOpt::edge()).unwrap();
            let _ = set_readiness.set_readiness(Ready::readable());
            drop(registration);
            drop(set_readiness);
            token_index += 1;

            thread::park_timeout(Duration::from_millis(0));
            let _ = poll.poll(&mut events, None).unwrap();
        }
    }
}
<|endoftext|>


<commit_before>extern crate cargotest;
extern crate hamcrest;
use cargotest::support::{project, execs, basic_bin_manifest};
use hamcrest::{assert_that};

#[test]
fn alias_incorrect_config_type() {
    let p = project("foo")
        .file("Cargo.toml", &basic_bin_manifest("foo"))
        .file("src/main.rs", r#"
            fn main() {
        }"#)
        .file(".cargo/config",r#"
            [alias]
            b-cargo-test = 5
        "#);;

    assert_that(p.cargo_process("b-cargo-test").arg("-v"),
                execs().with_status(101).
                with_stderr_contains("[ERROR] invalid configuration \
for key `alias.b-cargo-test`
expected a list, but found a integer for [..]"));
}

#[test]
fn alias_default_config_overrides_config() {
    let p = project("foo")
        .file("Cargo.toml", &basic_bin_manifest("foo"))
        .file("src/main.rs", r#"
            fn main() {
        }"#)
        .file(".cargo/config",r#"
            [alias]
            b = "not_build"
        "#);;

    assert_that(p.cargo_process("b").arg("-v"),
                execs().with_status(0).
                with_stderr_contains("[COMPILING] foo v0.5.0 [..]"));
}

#[test]
fn alias_config() {
    let p = project("foo")
        .file("Cargo.toml", &basic_bin_manifest("foo"))
        .file("src/main.rs", r#"
            fn main() {
        }"#)
        .file(".cargo/config",r#"
            [alias]
            b-cargo-test = "build"
        "#);;

    assert_that(p.cargo_process("b-cargo-test").arg("-v"),
                execs().with_status(0).
                with_stderr_contains("[COMPILING] foo v0.5.0 [..]
[RUNNING] `rustc --crate-name foo [..]"));
}

#[test]
fn alias_list_test() {
    let p = project("foo")
        .file("Cargo.toml", &basic_bin_manifest("foo"))
        .file("src/main.rs", r#"
            fn main() {
         }"#)
        .file(".cargo/config",r#"
            [alias]
            b-cargo-test = ["build", "--release"]
         "#);;

    assert_that(p.cargo_process("b-cargo-test").arg("-v"),
                execs().with_status(0).
                with_stderr_contains("[COMPILING] foo v0.5.0 [..]").
                with_stderr_contains("[RUNNING] `rustc --crate-name [..]")
                );
}

#[test]
fn alias_with_flags_config() {
    let p = project("foo")
        .file("Cargo.toml", &basic_bin_manifest("foo"))
        .file("src/main.rs", r#"
            fn main() {
         }"#)
        .file(".cargo/config",r#"
            [alias]
            b-cargo-test = "build --release"
         "#);;

    assert_that(p.cargo_process("b-cargo-test").arg("-v"),
                execs().with_status(0).
                with_stderr_contains("[COMPILING] foo v0.5.0 [..]").
                with_stderr_contains("[RUNNING] `rustc --crate-name foo [..]")
                );
}

#[test]
fn cant_shadow_builtin() {
    let p = project("foo")
        .file("Cargo.toml", &basic_bin_manifest("foo"))
        .file("src/main.rs", r#"
            fn main() {
         }"#)
        .file(".cargo/config",r#"
            [alias]
            build = "fetch"
         "#);;

    assert_that(p.cargo_process("build"),
                execs().with_status(0)
                       .with_stderr("\
[COMPILING] foo v0.5.0 ([..])
[FINISHED] dev [unoptimized + debuginfo] target(s) in [..]
"));
}
<commit_msg>Remove superfluous `;;` sequences<commit_after>extern crate cargotest;
extern crate hamcrest;
use cargotest::support::{project, execs, basic_bin_manifest};
use hamcrest::{assert_that};

#[test]
fn alias_incorrect_config_type() {
    let p = project("foo")
        .file("Cargo.toml", &basic_bin_manifest("foo"))
        .file("src/main.rs", r#"
            fn main() {
        }"#)
        .file(".cargo/config",r#"
            [alias]
            b-cargo-test = 5
        "#);

    assert_that(p.cargo_process("b-cargo-test").arg("-v"),
                execs().with_status(101).
                with_stderr_contains("[ERROR] invalid configuration \
for key `alias.b-cargo-test`
expected a list, but found a integer for [..]"));
}

#[test]
fn alias_default_config_overrides_config() {
    let p = project("foo")
        .file("Cargo.toml", &basic_bin_manifest("foo"))
        .file("src/main.rs", r#"
            fn main() {
        }"#)
        .file(".cargo/config",r#"
            [alias]
            b = "not_build"
        "#);

    assert_that(p.cargo_process("b").arg("-v"),
                execs().with_status(0).
                with_stderr_contains("[COMPILING] foo v0.5.0 [..]"));
}

#[test]
fn alias_config() {
    let p = project("foo")
        .file("Cargo.toml", &basic_bin_manifest("foo"))
        .file("src/main.rs", r#"
            fn main() {
        }"#)
        .file(".cargo/config",r#"
            [alias]
            b-cargo-test = "build"
        "#);

    assert_that(p.cargo_process("b-cargo-test").arg("-v"),
                execs().with_status(0).
                with_stderr_contains("[COMPILING] foo v0.5.0 [..]
[RUNNING] `rustc --crate-name foo [..]"));
}

#[test]
fn alias_list_test() {
    let p = project("foo")
        .file("Cargo.toml", &basic_bin_manifest("foo"))
        .file("src/main.rs", r#"
            fn main() {
         }"#)
        .file(".cargo/config",r#"
            [alias]
            b-cargo-test = ["build", "--release"]
         "#);

    assert_that(p.cargo_process("b-cargo-test").arg("-v"),
                execs().with_status(0).
                with_stderr_contains("[COMPILING] foo v0.5.0 [..]").
                with_stderr_contains("[RUNNING] `rustc --crate-name [..]")
                );
}

#[test]
fn alias_with_flags_config() {
    let p = project("foo")
        .file("Cargo.toml", &basic_bin_manifest("foo"))
        .file("src/main.rs", r#"
            fn main() {
         }"#)
        .file(".cargo/config",r#"
            [alias]
            b-cargo-test = "build --release"
         "#);

    assert_that(p.cargo_process("b-cargo-test").arg("-v"),
                execs().with_status(0).
                with_stderr_contains("[COMPILING] foo v0.5.0 [..]").
                with_stderr_contains("[RUNNING] `rustc --crate-name foo [..]")
                );
}

#[test]
fn cant_shadow_builtin() {
    let p = project("foo")
        .file("Cargo.toml", &basic_bin_manifest("foo"))
        .file("src/main.rs", r#"
            fn main() {
         }"#)
        .file(".cargo/config",r#"
            [alias]
            build = "fetch"
         "#);

    assert_that(p.cargo_process("build"),
                execs().with_status(0)
                       .with_stderr("\
[COMPILING] foo v0.5.0 ([..])
[FINISHED] dev [unoptimized + debuginfo] target(s) in [..]
"));
}
<|endoftext|>


<commit_before>#[cfg(feature = "ndarray_volumes")]
extern crate ndarray;
#[cfg(feature = "ndarray_volumes")]
extern crate nifti;
#[cfg(feature = "ndarray_volumes")]
extern crate tempfile;

#[cfg(feature = "ndarray_volumes")]
mod tests {
    use std::{
        fs::File,
        io::Read,
        ops::{Add, Mul},
        path::{Path, PathBuf},
    };

    use ndarray::{Array, Array2, Axis, Ix2, IxDyn, ShapeBuilder};
    use tempfile::tempdir;

    use nifti::{
        header::MAGIC_CODE_NIP1,
        object::NiftiObject,
        writer::{write_nifti, write_rgb_nifti},
        InMemNiftiObject, IntoNdArray, NiftiHeader, NiftiType,
    };

    fn get_temporary_path(ext: &str) -> PathBuf {
        let dir = tempdir().unwrap();
        let mut path = dir.into_path();
        if !ext.is_empty() {
            path.push(ext);
        }
        path
    }

    pub fn generate_nifti_header(
        dim: [u16; 8],
        scl_slope: f32,
        scl_inter: f32,
        datatype: i16,
    ) -> NiftiHeader {
        NiftiHeader {
            dim,
            datatype,
            bitpix: (NiftiType::Float32.size_of() * 8) as i16,
            magic: *MAGIC_CODE_NIP1,
            scl_slope,
            scl_inter,
            ..NiftiHeader::default()
        }
    }

    fn read_2d_image<P>(path: P) -> Array2<f32>
    where
        P: AsRef<Path>,
    {
        let nifti_object = InMemNiftiObject::from_file(path).expect("Nifti file is unreadable.");
        let volume = nifti_object.into_volume();
        let dyn_data = volume.into_ndarray().unwrap();
        dyn_data.into_dimensionality::<Ix2>().unwrap()
    }

    fn test_write_read(arr: &Array<f32, IxDyn>, path: &str) {
        let path = get_temporary_path(path);
        let mut dim = [1; 8];
        dim[0] = arr.ndim() as u16;
        for (i, s) in arr.shape().iter().enumerate() {
            dim[i + 1] = *s as u16;
        }
        let header = generate_nifti_header(dim, 1.0, 0.0, 16);
        write_nifti(&path, &arr, Some(&header)).unwrap();

        let read_nifti = read_2d_image(path);
        assert!(read_nifti.all_close(&arr, 1e-10));
    }

    fn f_order_array() -> Array<f32, IxDyn> {
        let dim = vec![4, 4];
        let vec = (0..16).map(|x| x as f32).collect();
        Array::from_shape_vec(IxDyn(&dim).f(), vec).unwrap()
    }

    fn c_order_array() -> Array<f32, IxDyn> {
        let dim = vec![4, 4];
        let vec = (0..16).map(|x| x as f32).collect();
        Array::from_shape_vec(IxDyn(&dim), vec).unwrap()
    }

    #[test]
    fn test_fortran_writing() {
        // Test .nii
        let arr = f_order_array();
        test_write_read(&arr, "test.nii");
        let mut arr = f_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii");

        // Test .nii.gz
        let arr = f_order_array();
        test_write_read(&arr, "test.nii.gz");
        let mut arr = f_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii.gz");
    }

    #[test]
    fn test_c_writing() {
        // Test .nii
        let arr = c_order_array();
        test_write_read(&arr, "test.nii");

        let mut arr = c_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii");

        // Test .nii.gz
        let arr = c_order_array();
        test_write_read(&arr, "test.nii.gz");

        let mut arr = c_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii.gz");
    }

    #[test]
    fn test_header_slope_inter() {
        let arr = f_order_array();
        let slope = 2.2;
        let inter = 101.1;

        let path = get_temporary_path("test_slope_inter.nii");
        let mut dim = [1; 8];
        dim[0] = arr.ndim() as u16;
        for (i, s) in arr.shape().iter().enumerate() {
            dim[i + 1] = *s as u16;
        }
        let header = generate_nifti_header(dim, slope, inter, 16);
        let transformed_data = arr.mul(slope).add(inter);
        write_nifti(&path, &transformed_data, Some(&header)).unwrap();

        let read_nifti = read_2d_image(path);
        assert!(read_nifti.all_close(&transformed_data, 1e-10));
    }

    #[test]
    fn test_write_3d_rgb() {
        let mut data = Array::from_elem((3, 3, 3), [0u8, 0u8, 0u8]);
        data[(0, 0, 0)] = [55, 55, 0];
        data[(0, 0, 1)] = [55, 0, 55];
        data[(0, 1, 0)] = [0, 55, 55];

        let path = get_temporary_path("rgb.nii");
        write_rgb_nifti(&path, &data, None).unwrap();

        // Until we are able to read RGB images, we simply compare the bytes of the newly created
        // image to the bytes of the prepared 3D RGB image in ressources/rgb/.
        let mut rgb_bytes = vec![];
        let _ = File::open(path)
            .unwrap()
            .read_to_end(&mut rgb_bytes)
            .unwrap();
        let mut gt_bytes = vec![];
        let _ = File::open("resources/rgb/3D.nii")
            .unwrap()
            .read_to_end(&mut gt_bytes)
            .unwrap();
        assert_eq!(rgb_bytes, gt_bytes);
    }

    #[test]
    fn test_write_4d_rgb() {
        let mut data = Array::from_elem((3, 3, 3, 2), [0u8, 0u8, 0u8]);
        data[(0, 0, 0, 0)] = [55, 55, 0];
        data[(0, 0, 1, 0)] = [55, 0, 55];
        data[(0, 1, 0, 0)] = [0, 55, 55];
        data[(0, 0, 0, 1)] = [55, 55, 0];
        data[(0, 1, 0, 1)] = [55, 0, 55];
        data[(1, 0, 0, 1)] = [0, 55, 55];

        let path = get_temporary_path("rgb.nii");
        write_rgb_nifti(&path, &data, None).unwrap();

        // Until we are able to read RGB images, we simply compare the bytes of the newly created
        // image to the bytes of the prepared 4D RGB image in ressources/rgb/.
        let mut rgb_bytes = vec![];
        let _ = File::open(path)
            .unwrap()
            .read_to_end(&mut rgb_bytes)
            .unwrap();
        let mut gt_bytes = vec![];
        let _ = File::open("resources/rgb/4D.nii")
            .unwrap()
            .read_to_end(&mut gt_bytes)
            .unwrap();
        assert_eq!(rgb_bytes, gt_bytes);
    }
}
<commit_msg>Clean writing tests<commit_after>#[cfg(feature = "ndarray_volumes")]
extern crate ndarray;
#[cfg(feature = "ndarray_volumes")]
extern crate nifti;
#[cfg(feature = "ndarray_volumes")]
extern crate tempfile;

#[cfg(feature = "ndarray_volumes")]
mod tests {
    use std::{
        fs::File,
        io::Read,
        ops::{Add, Mul},
        path::{Path, PathBuf},
    };

    use ndarray::{Array, Array2, Axis, Dimension, IxDyn, ShapeBuilder};
    use tempfile::tempdir;

    use nifti::{
        header::MAGIC_CODE_NIP1,
        object::NiftiObject,
        writer::{write_nifti, write_rgb_nifti},
        InMemNiftiObject, IntoNdArray, NiftiHeader, NiftiType,
    };

    fn get_temporary_path(ext: &str) -> PathBuf {
        let dir = tempdir().unwrap();
        let mut path = dir.into_path();
        if !ext.is_empty() {
            path.push(ext);
        }
        path
    }

    pub fn generate_nifti_header(
        dim: [u16; 8],
        scl_slope: f32,
        scl_inter: f32,
        datatype: i16,
    ) -> NiftiHeader {
        NiftiHeader {
            dim,
            datatype,
            bitpix: (NiftiType::Float32.size_of() * 8) as i16,
            magic: *MAGIC_CODE_NIP1,
            scl_slope,
            scl_inter,
            ..NiftiHeader::default()
        }
    }

    fn read_as_ndarray<P, D>(path: P) -> Array<f32, D>
    where
        P: AsRef<Path>,
        D: Dimension,
    {
        let nifti_object = InMemNiftiObject::from_file(path).expect("Nifti file is unreadable.");
        let volume = nifti_object.into_volume();
        let dyn_data = volume.into_ndarray().unwrap();
        dyn_data.into_dimensionality::<D>().unwrap()
    }

    fn test_write_read(arr: &Array<f32, IxDyn>, path: &str) {
        let path = get_temporary_path(path);
        let mut dim = [1; 8];
        dim[0] = arr.ndim() as u16;
        for (i, s) in arr.shape().iter().enumerate() {
            dim[i + 1] = *s as u16;
        }
        let header = generate_nifti_header(dim, 1.0, 0.0, 16);
        write_nifti(&path, &arr, Some(&header)).unwrap();

        let read_nifti: Array2<f32> = read_as_ndarray(path);
        assert!(read_nifti.all_close(&arr, 1e-10));
    }

    fn f_order_array() -> Array<f32, IxDyn> {
        let dim = vec![4, 4];
        let vec = (0..16).map(|x| x as f32).collect();
        Array::from_shape_vec(IxDyn(&dim).f(), vec).unwrap()
    }

    fn c_order_array() -> Array<f32, IxDyn> {
        let dim = vec![4, 4];
        let vec = (0..16).map(|x| x as f32).collect();
        Array::from_shape_vec(IxDyn(&dim), vec).unwrap()
    }

    #[test]
    fn fortran_writing() {
        // Test .nii
        let arr = f_order_array();
        test_write_read(&arr, "test.nii");
        let mut arr = f_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii");

        // Test .nii.gz
        let arr = f_order_array();
        test_write_read(&arr, "test.nii.gz");
        let mut arr = f_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii.gz");
    }

    #[test]
    fn c_writing() {
        // Test .nii
        let arr = c_order_array();
        test_write_read(&arr, "test.nii");

        let mut arr = c_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii");

        // Test .nii.gz
        let arr = c_order_array();
        test_write_read(&arr, "test.nii.gz");

        let mut arr = c_order_array();
        arr.invert_axis(Axis(1));
        test_write_read(&arr, "test_non_contiguous.nii.gz");
    }

    #[test]
    fn header_slope_inter() {
        let arr = f_order_array();
        let slope = 2.2;
        let inter = 101.1;

        let path = get_temporary_path("test_slope_inter.nii");
        let mut dim = [1; 8];
        dim[0] = arr.ndim() as u16;
        for (i, s) in arr.shape().iter().enumerate() {
            dim[i + 1] = *s as u16;
        }
        let header = generate_nifti_header(dim, slope, inter, 16);
        let transformed_data = arr.mul(slope).add(inter);
        write_nifti(&path, &transformed_data, Some(&header)).unwrap();

        let read_nifti: Array2<f32> = read_as_ndarray(path);
        assert!(read_nifti.all_close(&transformed_data, 1e-10));
    }

    #[test]
    fn write_3d_rgb() {
        let mut data = Array::from_elem((3, 3, 3), [0u8, 0u8, 0u8]);
        data[(0, 0, 0)] = [55, 55, 0];
        data[(0, 0, 1)] = [55, 0, 55];
        data[(0, 1, 0)] = [0, 55, 55];

        let path = get_temporary_path("rgb.nii");
        write_rgb_nifti(&path, &data, None).unwrap();

        // Until we are able to read RGB images, we simply compare the bytes of the newly created
        // image to the bytes of the prepared 3D RGB image in ressources/rgb/.
        let mut rgb_bytes = vec![];
        let _ = File::open(path)
            .unwrap()
            .read_to_end(&mut rgb_bytes)
            .unwrap();
        let mut gt_bytes = vec![];
        let _ = File::open("resources/rgb/3D.nii")
            .unwrap()
            .read_to_end(&mut gt_bytes)
            .unwrap();
        assert_eq!(rgb_bytes, gt_bytes);
    }

    #[test]
    fn write_4d_rgb() {
        let mut data = Array::from_elem((3, 3, 3, 2), [0u8, 0u8, 0u8]);
        data[(0, 0, 0, 0)] = [55, 55, 0];
        data[(0, 0, 1, 0)] = [55, 0, 55];
        data[(0, 1, 0, 0)] = [0, 55, 55];
        data[(0, 0, 0, 1)] = [55, 55, 0];
        data[(0, 1, 0, 1)] = [55, 0, 55];
        data[(1, 0, 0, 1)] = [0, 55, 55];

        let path = get_temporary_path("rgb.nii");
        write_rgb_nifti(&path, &data, None).unwrap();

        // Until we are able to read RGB images, we simply compare the bytes of the newly created
        // image to the bytes of the prepared 4D RGB image in ressources/rgb/.
        let mut rgb_bytes = vec![];
        let _ = File::open(path)
            .unwrap()
            .read_to_end(&mut rgb_bytes)
            .unwrap();
        let mut gt_bytes = vec![];
        let _ = File::open("resources/rgb/4D.nii")
            .unwrap()
            .read_to_end(&mut gt_bytes)
            .unwrap();
        assert_eq!(rgb_bytes, gt_bytes);
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::fs::File;
use std::ops::Drop;
use std::path::PathBuf;
use std::result::Result as RResult;
use std::sync::Arc;
use std::sync::RwLock;

pub use store::Store;
pub use store::Result;
pub use entry::Entry;
pub use error::StoreError;

pub struct FSStore {
    location: PathBuf,

    /**
     * Internal Path->File cache map
     *
     * Caches the files, so they remain flock()ed
     *
     * Could be optimized for a threadsafe HashMap
     */
    cache: Arc<RwLock<HashMap<PathBuf, File>>>,
}

impl FSStore {

    pub fn new(location: PathBuf) -> FSStore {
        FSStore {
            location: location,
            cache: Arc::new(RwLock::new(HashMap::new())),
        }
    }
}

impl Store for FSStore {

    fn location(&self) -> &PathBuf {
        &self.location
    }

    fn create(&self, entry: Entry) -> Result<()> {
        unimplemented!()
    }

    fn read(&self, path: PathBuf) -> Result<Arc<RwLock<Entry>>> {
        unimplemented!()
    }

    fn update(&self, entry: Arc<RwLock<Entry>>) -> Result<()> {
        unimplemented!()
    }

    fn delete(&self, path: PathBuf) -> Result<()> {
        unimplemented!()
    }

}

impl Drop for FSStore {

    /**
     * Unlock all files on drop
     *
     * TODO: Error message when file cannot be unlocked?
     */
    fn drop(&mut self) {
        use std::ops::DerefMut;
        use fs2::FileExt;

        let cache = self.cache.clone();
        cache.write().unwrap().deref_mut().iter().map(|(_, f)| f.unlock());
    }

}

<commit_msg>Add missing method retrieve_copy<commit_after>use std::collections::HashMap;
use std::fs::File;
use std::ops::Drop;
use std::path::PathBuf;
use std::result::Result as RResult;
use std::sync::Arc;
use std::sync::RwLock;

pub use store::Store;
pub use store::Result;
pub use entry::Entry;
pub use error::StoreError;

pub struct FSStore {
    location: PathBuf,

    /**
     * Internal Path->File cache map
     *
     * Caches the files, so they remain flock()ed
     *
     * Could be optimized for a threadsafe HashMap
     */
    cache: Arc<RwLock<HashMap<PathBuf, File>>>,
}

impl FSStore {

    pub fn new(location: PathBuf) -> FSStore {
        FSStore {
            location: location,
            cache: Arc::new(RwLock::new(HashMap::new())),
        }
    }
}

impl Store for FSStore {

    fn location(&self) -> &PathBuf {
        &self.location
    }

    fn create(&self, entry: Entry) -> Result<()> {
        unimplemented!()
    }

    fn read(&self, path: PathBuf) -> Result<Arc<RwLock<Entry>>> {
        unimplemented!()
    }

    fn retrieve_copy(&self, id : String) -> Result<Entry> {
        unimplemented!()
    }

    fn update(&self, entry: Arc<RwLock<Entry>>) -> Result<()> {
        unimplemented!()
    }

    fn delete(&self, path: PathBuf) -> Result<()> {
        unimplemented!()
    }

}

impl Drop for FSStore {

    /**
     * Unlock all files on drop
     *
     * TODO: Error message when file cannot be unlocked?
     */
    fn drop(&mut self) {
        use std::ops::DerefMut;
        use fs2::FileExt;

        let cache = self.cache.clone();
        cache.write().unwrap().deref_mut().iter().map(|(_, f)| f.unlock());
    }

}

<|endoftext|>


<commit_before>
/**
 Elements that may be contained by a quadtree node.
 Either child nodes, a single rect, or nothing.
 */
enum Elements {
	/// Children are top left, top right, bottom right, and bottom left, respectively.
	Children(~QuadTree, ~QuadTree, ~QuadTree, ~QuadTree),
	/// A single rectangle.
	Member(Rect),
	/// Nothing.
	None
}

pub struct Point {
	x: f32,
	y: f32,
}

pub struct Size {
	width: f32,
	height: f32,
}

pub struct Rect {
	origin: Point,
	size: Size,
}

/**
 A quadtree node that can contain either one rectangle,
 or exactly four child nodes.
 */
pub struct QuadTree {
	rect: Rect,
	elements: Elements,
}

impl Rect {
	fn new(origin: Point, size: Size) -> Rect {
		Rect { origin: origin, size: size }
	}

	fn minXRect<'a>(rect1: &'a Rect, rect2: &'a Rect) -> (&'a Rect, &'a Rect) {
		// Find which the rect has an origin farther to the left.
		if (rect1.minX() <= rect2.minX()) {
			(rect1, rect2)
		} else {
			(rect2, rect1)
		}
	}

	/**
	 Check if this rect entirely contains another rect.
	 */
	fn contains(&self, rect: &Rect) -> bool {
		// Find which the rect has an origin farther to the left.
		let (minXRect, otherRect) = Rect::minXRect(self, rect);

		// If the rectangles don't intersect, one cannot be contained in the other.
		self.intersects(rect)
		&& (minXRect.maxX() >= otherRect.maxX() && minXRect.maxY() >= otherRect.maxY())
	}

	/**
	 Check if this rect and another rect intersect.
	 */
	fn intersects(&self, rect: &Rect) -> bool {
		let (minXRect, otherRect) = Rect::minXRect(self, rect);

		let intersects: bool = ((minXRect.maxX() >= otherRect.minX())
			&& ((minXRect.minY() >= otherRect.minY() && minXRect.minY() <= otherRect.maxY())
				|| (minXRect.minY() <= otherRect.minY() && minXRect.maxY() >= otherRect.minY())));

		intersects
	}

	fn maxX(&self) -> f32 {
		self.origin.x + self.size.width
	}

	fn maxY(&self) -> f32 {
		self.origin.y + self.size.height
	}

	fn minX(&self) -> f32 {
		self.origin.x
	}

	fn minY(&self) -> f32 {
		self.origin.y
	}
}

impl QuadTree {
	/**
	 Create a quadtree with a square root node with the same origin and size
	 as the passed in rect.
	 */
	fn newAutosized(rect: Rect) -> QuadTree {
		QuadTree { rect: rect, elements: Member(rect), }
	}

	/**
	 Create a quadtree with a root node with the given origin, size, and member rectangle.
	 If `insertRect` is not entirely contained within the rectangle defined by `origin` and `size`,
	 higher level nodes will be created to form a quadtree with a child of the rect specified
	 and other children such that the rect is in the tree.
	 */
	fn new(origin: Point, size: Size, insertRect: Rect) -> QuadTree {
		let nodeRect = Rect { origin: origin, size: size };
		let node = QuadTree { rect: nodeRect, elements: None };

		node.insertRect(insertRect)
	}

	fn insertRect(self, toInsert: Rect) -> QuadTree {
		fail!("insertRect not implemented!");

		self
	}
}<commit_msg>Turn a comment into a doc comment.<commit_after>
/**
 Elements that may be contained by a quadtree node.
 Either child nodes, a single rect, or nothing.
 */
enum Elements {
	/// Children are top left, top right, bottom right, and bottom left, respectively.
	Children(~QuadTree, ~QuadTree, ~QuadTree, ~QuadTree),
	/// A single rectangle.
	Member(Rect),
	/// Nothing.
	None
}

pub struct Point {
	x: f32,
	y: f32,
}

pub struct Size {
	width: f32,
	height: f32,
}

pub struct Rect {
	origin: Point,
	size: Size,
}

/**
 A quadtree node that can contain either one rectangle,
 or exactly four child nodes.
 */
pub struct QuadTree {
	rect: Rect,
	elements: Elements,
}

impl Rect {
	fn new(origin: Point, size: Size) -> Rect {
		Rect { origin: origin, size: size }
	}

	/// Find which the rect has an origin farther to the left.
	fn minXRect<'a>(rect1: &'a Rect, rect2: &'a Rect) -> (&'a Rect, &'a Rect) {
		if (rect1.minX() <= rect2.minX()) {
			(rect1, rect2)
		} else {
			(rect2, rect1)
		}
	}

	/**
	 Check if this rect entirely contains another rect.
	 */
	fn contains(&self, rect: &Rect) -> bool {
		// Find which the rect has an origin farther to the left.
		let (minXRect, otherRect) = Rect::minXRect(self, rect);

		// If the rectangles don't intersect, one cannot be contained in the other.
		self.intersects(rect)
		&& (minXRect.maxX() >= otherRect.maxX() && minXRect.maxY() >= otherRect.maxY())
	}

	/**
	 Check if this rect and another rect intersect.
	 */
	fn intersects(&self, rect: &Rect) -> bool {
		let (minXRect, otherRect) = Rect::minXRect(self, rect);

		let intersects: bool = ((minXRect.maxX() >= otherRect.minX())
			&& ((minXRect.minY() >= otherRect.minY() && minXRect.minY() <= otherRect.maxY())
				|| (minXRect.minY() <= otherRect.minY() && minXRect.maxY() >= otherRect.minY())));

		intersects
	}

	fn maxX(&self) -> f32 {
		self.origin.x + self.size.width
	}

	fn maxY(&self) -> f32 {
		self.origin.y + self.size.height
	}

	fn minX(&self) -> f32 {
		self.origin.x
	}

	fn minY(&self) -> f32 {
		self.origin.y
	}
}

impl QuadTree {
	/**
	 Create a quadtree with a square root node with the same origin and size
	 as the passed in rect.
	 */
	fn newAutosized(rect: Rect) -> QuadTree {
		QuadTree { rect: rect, elements: Member(rect), }
	}

	/**
	 Create a quadtree with a root node with the given origin, size, and member rectangle.
	 If `insertRect` is not entirely contained within the rectangle defined by `origin` and `size`,
	 higher level nodes will be created to form a quadtree with a child of the rect specified
	 and other children such that the rect is in the tree.
	 */
	fn new(origin: Point, size: Size, insertRect: Rect) -> QuadTree {
		let nodeRect = Rect { origin: origin, size: size };
		let node = QuadTree { rect: nodeRect, elements: None };

		node.insertRect(insertRect)
	}

	fn insertRect(self, toInsert: Rect) -> QuadTree {
		fail!("insertRect not implemented!");

		self
	}
}<|endoftext|>


<commit_before>// https://doc.rust-lang.org/nightly/std/ops/trait.Try.html
// https://github.com/rust-lang/rfcs/blob/master/text/3058-try-trait-v2.md
// https://github.com/rust-lang/rust/issues/84277#issuecomment-905821708
// https://blog.rust-lang.org/2020/08/27/Rust-1.46.0.html#track_caller

use fomat_macros::fomat;
use std::fmt;
use std::ops::{ControlFlow, FromResidual, Try};
use std::panic::Location;
use std::process::{ExitCode, Termination};
use super::filename;

#[macro_export]
macro_rules! fail {($($args: tt)+) => (return $crate::re::Re::fail (fomat! ($($args)+)))}

/// fail with a small message
#[cfg(feature = "fomat-macros")]
#[macro_export] macro_rules! ifail {
  ($($args: tt)+) => (return Re::fail (ifomat! ($($args)+)))}

#[derive(Debug, PartialEq)]
#[must_use = "this `Re` may be an `Err` variant, which should be handled"]
pub enum Re<T> {Ok (T), Err (String)}

impl<T> Termination for Re<T> {
  fn report (self) -> ExitCode {
    match self {
      Re::Ok (_) => ExitCode::SUCCESS,
      Re::Err (err) => {
        eprintln! ("{}", err);
        use std::hash::{Hash, Hasher};
        let mut hasher = std::collections::hash_map::DefaultHasher::new();
        err.hash (&mut hasher);
        ExitCode::from (1 | (hasher.finish() & 0xFF) as u8)}}}}

impl<T> Try for Re<T> {
  type Output = T;
  type Residual = Re<!>;

  #[inline]
  fn from_output (c: T) -> Self {Re::Ok (c)}

  #[inline]
  fn branch (self) -> ControlFlow<Self::Residual, T> {
    match self {
      Re::Ok (c) => ControlFlow::Continue (c),
      Re::Err (e) => ControlFlow::Break (Re::Err (e))}}}

impl<T> FromResidual<Re<!>> for Re<T> {
  #[track_caller]
  fn from_residual (x: Re<!>) -> Self {
    let err = match x {Re::Ok (_) => unreachable!(), Re::Err (e) => e};
    let loc = Location::caller();
    let err = fomat! ((filename (loc.file())) ':' (loc.line()) "] " (err));
    Re::Err (err)}}

/// Allows `Re<_>` errors to chain into `Result<_, String>` errors
/// which can be useful with libraries using `Result`, `OnceCell::get_or_try_init` for example
impl<T, O> FromResidual<Re<O>> for Result<T, String> {
  #[track_caller]
  fn from_residual (x: Re<O>) -> Self {
    let err = match x {Re::Ok (_) => unreachable!(), Re::Err (e) => e};
    let loc = Location::caller();
    let err = fomat! ((filename (loc.file())) ':' (loc.line()) "] " (err));
    Result::Err (err)}}

impl<T, O, E> FromResidual<Result<O, E>> for Re<T> where E: fmt::Display {
  #[track_caller]
  fn from_residual (x: Result<O, E>) -> Self {
    let err = match x {Result::Ok (_) => unreachable!(), Result::Err (e) => e};
    let loc = Location::caller();
    let err = fomat! ((filename (loc.file())) ':' (loc.line()) "] " (err));
    Re::Err (err)}}

impl<T, O> FromResidual<Option<O>> for Re<T> {
  #[track_caller]
  fn from_residual (_: Option<O>) -> Self {
    let loc = Location::caller();
    let err = fomat! ((filename (loc.file())) ':' (loc.line()) "] Option is None");
    Re::Err (err)}}

impl<T> Re<T> {
  #[track_caller]
  pub fn fail<E: fmt::Display> (emsg: E) -> Re<T> {
    let loc = Location::caller();
    let err = fomat! ((filename (loc.file())) ':' (loc.line()) "] " (emsg));
    Re::Err (err)}

  #[inline]
  pub fn err (self) -> Option<String> {
    match self {Re::Ok (_) => None, Re::Err (e) => Some (e)}}

  #[inline]
  #[track_caller]
  pub fn expect (self, msg: &str) -> T {
    match self {Re::Ok (k) => k, Re::Err (err) => panic! ("{}: {:?}", msg, err)}}

  #[inline]
  #[track_caller]
  pub fn unwrap (self) -> T {
    match self {Re::Ok (k) => k, Re::Err (err) => panic! ("called `Re::unwrap()` on an `Err` value: {:?}", err)}}

  #[inline]
  pub fn unwrap_or (self, default: T) -> T {
    match self {Re::Ok (k) => k, Re::Err (_) => default}}

  #[inline]
  pub fn map<U, F: FnOnce (T) -> U> (self, op: F) -> Re<U> {
    match self {Re::Ok (k) => Re::Ok (op (k)), Re::Err (e) => Re::Err (e)}}}

#[cfg(all(test, feature = "nightly"))] mod test {
  extern crate test;

  use super::*;

  #[test] fn opt() {
    fn foo() -> Re<()> {
      let bar: Option<String> = None;
      bar?;
      Re::Ok(())}
    assert_eq! (foo(), Re::Err ("re:119] Option is None".into()));
    //assert_eq! (foo().report(), 521090)
  }

  #[test] fn re2result() {
    fn foo() -> Result<(), String> {
      let bar: Re<()> = Re::fail ("ups");
      bar?;
      Result::Ok(())}
    assert_eq! (foo(), Result::Err ("re:128] re:127] ups".into()));
    // #![feature(process_exitcode_internals)]
    //assert_eq! (foo().report().to_i32(), 1)
  }

  // cargo bench --features nightly,re

  fn foobar (succ: bool) -> Re<bool> {
    if test::black_box (succ) {Re::Ok (true)} else {Re::fail ("!succ")}}

  fn bang (succ: bool) -> Re<()> {
    let b = test::black_box (foobar (succ)?);
    assert_eq! (true, b);
    Re::Ok(())}

  #[bench] fn err (bm: &mut test::Bencher) {
    bm.iter (|| {
      let e = bang (false) .err().unwrap();
      assert! (e.starts_with ("re:"));
      assert! (e.ends_with ("] !succ"))})}

  #[bench] fn ok (bm: &mut test::Bencher) {
    bm.iter (|| {bang (true) .expect ("!bang")})}}
<commit_msg>Add Re::is_ok<commit_after>// https://doc.rust-lang.org/nightly/std/ops/trait.Try.html
// https://github.com/rust-lang/rfcs/blob/master/text/3058-try-trait-v2.md
// https://github.com/rust-lang/rust/issues/84277#issuecomment-905821708
// https://blog.rust-lang.org/2020/08/27/Rust-1.46.0.html#track_caller

use fomat_macros::fomat;
use std::fmt;
use std::ops::{ControlFlow, FromResidual, Try};
use std::panic::Location;
use std::process::{ExitCode, Termination};
use super::filename;

#[macro_export]
macro_rules! fail {($($args: tt)+) => (return $crate::re::Re::fail (fomat! ($($args)+)))}

/// fail with a small message
#[cfg(feature = "fomat-macros")]
#[macro_export] macro_rules! ifail {
  ($($args: tt)+) => (return Re::fail (ifomat! ($($args)+)))}

#[derive(Debug, PartialEq)]
#[must_use = "this `Re` may be an `Err` variant, which should be handled"]
pub enum Re<T> {Ok (T), Err (String)}

impl<T> Termination for Re<T> {
  fn report (self) -> ExitCode {
    match self {
      Re::Ok (_) => ExitCode::SUCCESS,
      Re::Err (err) => {
        eprintln! ("{}", err);
        use std::hash::{Hash, Hasher};
        let mut hasher = std::collections::hash_map::DefaultHasher::new();
        err.hash (&mut hasher);
        ExitCode::from (1 | (hasher.finish() & 0xFF) as u8)}}}}

impl<T> Try for Re<T> {
  type Output = T;
  type Residual = Re<!>;

  #[inline]
  fn from_output (c: T) -> Self {Re::Ok (c)}

  #[inline]
  fn branch (self) -> ControlFlow<Self::Residual, T> {
    match self {
      Re::Ok (c) => ControlFlow::Continue (c),
      Re::Err (e) => ControlFlow::Break (Re::Err (e))}}}

impl<T> FromResidual<Re<!>> for Re<T> {
  #[track_caller]
  fn from_residual (x: Re<!>) -> Self {
    let err = match x {Re::Ok (_) => unreachable!(), Re::Err (e) => e};
    let loc = Location::caller();
    let err = fomat! ((filename (loc.file())) ':' (loc.line()) "] " (err));
    Re::Err (err)}}

/// Allows `Re<_>` errors to chain into `Result<_, String>` errors
/// which can be useful with libraries using `Result`, `OnceCell::get_or_try_init` for example
impl<T, O> FromResidual<Re<O>> for Result<T, String> {
  #[track_caller]
  fn from_residual (x: Re<O>) -> Self {
    let err = match x {Re::Ok (_) => unreachable!(), Re::Err (e) => e};
    let loc = Location::caller();
    let err = fomat! ((filename (loc.file())) ':' (loc.line()) "] " (err));
    Result::Err (err)}}

impl<T, O, E> FromResidual<Result<O, E>> for Re<T> where E: fmt::Display {
  #[track_caller]
  fn from_residual (x: Result<O, E>) -> Self {
    let err = match x {Result::Ok (_) => unreachable!(), Result::Err (e) => e};
    let loc = Location::caller();
    let err = fomat! ((filename (loc.file())) ':' (loc.line()) "] " (err));
    Re::Err (err)}}

impl<T, O> FromResidual<Option<O>> for Re<T> {
  #[track_caller]
  fn from_residual (_: Option<O>) -> Self {
    let loc = Location::caller();
    let err = fomat! ((filename (loc.file())) ':' (loc.line()) "] Option is None");
    Re::Err (err)}}

impl<T> Re<T> {
  #[track_caller]
  pub fn fail<E: fmt::Display> (emsg: E) -> Re<T> {
    let loc = Location::caller();
    let err = fomat! ((filename (loc.file())) ':' (loc.line()) "] " (emsg));
    Re::Err (err)}

  #[inline]
  pub fn is_ok (self) -> bool {
    match self {Re::Ok (_) => true, _ => false}}

  #[inline]
  pub fn err (self) -> Option<String> {
    match self {Re::Ok (_) => None, Re::Err (e) => Some (e)}}

  #[inline]
  #[track_caller]
  pub fn expect (self, msg: &str) -> T {
    match self {Re::Ok (k) => k, Re::Err (err) => panic! ("{}: {:?}", msg, err)}}

  #[inline]
  #[track_caller]
  pub fn unwrap (self) -> T {
    match self {Re::Ok (k) => k, Re::Err (err) => panic! ("called `Re::unwrap()` on an `Err` value: {:?}", err)}}

  #[inline]
  pub fn unwrap_or (self, default: T) -> T {
    match self {Re::Ok (k) => k, Re::Err (_) => default}}

  #[inline]
  pub fn map<U, F: FnOnce (T) -> U> (self, op: F) -> Re<U> {
    match self {Re::Ok (k) => Re::Ok (op (k)), Re::Err (e) => Re::Err (e)}}}

#[cfg(all(test, feature = "nightly"))] mod test {
  extern crate test;

  use super::*;

  #[test] fn opt() {
    fn foo() -> Re<()> {
      let bar: Option<String> = None;
      bar?;
      Re::Ok(())}
    assert_eq! (foo(), Re::Err ("re:119] Option is None".into()));
    //assert_eq! (foo().report(), 521090)
  }

  #[test] fn re2result() {
    fn foo() -> Result<(), String> {
      let bar: Re<()> = Re::fail ("ups");
      bar?;
      Result::Ok(())}
    assert_eq! (foo(), Result::Err ("re:128] re:127] ups".into()));
    // #![feature(process_exitcode_internals)]
    //assert_eq! (foo().report().to_i32(), 1)
  }

  // cargo bench --features nightly,re

  fn foobar (succ: bool) -> Re<bool> {
    if test::black_box (succ) {Re::Ok (true)} else {Re::fail ("!succ")}}

  fn bang (succ: bool) -> Re<()> {
    let b = test::black_box (foobar (succ)?);
    assert_eq! (true, b);
    Re::Ok(())}

  #[bench] fn err (bm: &mut test::Bencher) {
    bm.iter (|| {
      let e = bang (false) .err().unwrap();
      assert! (e.starts_with ("re:"));
      assert! (e.ends_with ("] !succ"))})}

  #[bench] fn ok (bm: &mut test::Bencher) {
    bm.iter (|| {bang (true) .expect ("!bang")})}}
<|endoftext|>