<commit_before>extern crate libc;
use self::libc::*;

use low_level as ll;
use error::Error;
use device::Device;
use context::Context;
use event::{Event, EventStatus, EventLike};
use kernel::Kernel;
use mem::Mem;

use std::ptr;
use std::mem::{size_of, transmute};

pub struct CommandQueue {
    pub id: ll::CommandQueue,
}

pub struct ReadLock<'a, T: Sized + 'a> {
    #[allow(dead_code)]
    sl: &'a [T],
    evt: Event,
}

pub struct WriteLock<'a, T: Sized + 'a> {
    #[allow(dead_code)]
    sl: &'a mut [T],
    evt: Event,
}

impl<'a, T: Sized> Drop for ReadLock<'a, T> {
    fn drop(self: &mut Self) -> () {
        self.evt.wait().ok().expect(
            "error waiting to finish write for ReadLock");
    }
}

impl<'a, T: Sized> EventLike for ReadLock<'a, T> {
    fn status(self: &Self) -> Result<EventStatus, Error> {
        self.evt.status()
    }

    fn wait(self: &Self) -> Result<(), Error> {
        self.evt.wait()
    }
}

impl<'a, T: Sized> Drop for WriteLock<'a, T> {
    fn drop(self: &mut Self) -> () {
        self.evt.wait().ok().expect(
            "error waiting to finish write for ReadLock");
    }
}

impl<'a, T: Sized> EventLike for WriteLock<'a, T> {
    fn status(self: &Self) -> Result<EventStatus, Error> {
        self.evt.status()
    }

    fn wait(self: &Self) -> Result<(), Error> {
        self.evt.wait()
    }
}

impl CommandQueue {
    pub fn new(context: Context,
               device: Device) -> Result<CommandQueue, Error> {
        let props: c_ulong = if try!(device.out_of_order_supported()) {
            1
        } else {
            0
        };
        let id = unsafe {
            let mut err: i32 = 0;
            let id = ll::clCreateCommandQueue(context.id,
                                              device.id,
                                              props,
                                              &mut err);
            try!(Error::check(err));
            id
        };
        Ok(CommandQueue{ id: id })
    }

    fn get_info(self: &Self, param: u32) -> Result<*mut c_void, Error> {
        let mut tr: *mut c_void = ptr::null_mut();
        unsafe {
            try!(Error::check(ll::clGetCommandQueueInfo(self.id,
                                                   param,
                                                   size_of::<*mut c_void>() as size_t,
                                                   transmute(&mut tr),
                                                   ptr::null_mut())));
        }
        Ok(tr)
    }

    pub fn context(self: &Self) -> Result<Context, Error> {
        let val = try!(self.get_info(0x1090));
        let id = unsafe { transmute(val) };
        unsafe {
            try!(Error::check(ll::clRetainContext(id)));
        }
        Ok(Context{id: id})
    }

    pub fn device(self: &Self) -> Result<Device, Error> {
        let val = try!(self.get_info(0x1091));
        let id = unsafe { transmute(val) };
        Ok(Device{id: id})
    }

    fn alloc_ll_buffer(self: &Self,
                    flags: c_ulong,
                    size_bytes: usize) -> Result<ll::Mem, Error> {
        let context = try!(self.context());
        let id = unsafe {
            let mut err: i32 = 0;
            let id = ll::clCreateBuffer(context.id,
                                        flags,
                                        size_bytes as size_t,
                                        ptr::null_mut(),
                                        &mut err);
            try!(Error::check(err));
            id
        };
        Ok(id)
    }

    fn alloc_ll_buffer_with_data(self: &Self,
                                 flags: c_ulong,
                                 size_bytes: usize,
                                 data: *const c_void) -> Result<ll::Mem, Error> {
        let context = try!(self.context());
        let id = unsafe {
            let mut err: i32 = 0;
            let id = ll::clCreateBuffer(context.id,
                                        flags | (1 << 5),
                                        size_bytes as size_t,
                                        data,
                                        &mut err);
            try!(Error::check(err));
            id
        };
        Ok(id)
    }

    fn write_buffer_raw<T: Sized>(self: &Self,
                       buf: &mut Mem,
                       offset: usize,
                       size: usize,
                       data: &[T]) -> Result<Event, Error> {
        let mut event_id: ll::Event = ptr::null_mut();
        let host_ptr = &data[0] as *const T;
        let wait_list: Vec<ll::Event> = buf.write_events();
        let wait_size = wait_list.len() as u32;

        let wait_list_ptr: *const ll::Event = match wait_size {
            0 => ptr::null(),
            _ => &wait_list[0],
        };

        unsafe {
            try!(Error::check(ll::clEnqueueWriteBuffer(self.id,
                                                       buf.id,
                                                       0, // don't block
                                                       offset as size_t,
                                                       size as size_t,
                                                       transmute(host_ptr),
                                                       wait_size,
                                                       wait_list_ptr,
                                                       &mut event_id)));
        }

        let event = Event{id: event_id};
        buf.register_write(event.clone());
        Ok(event)
    }

    fn read_buffer_raw<T: Sized>(self: &Self,
                                 buf: &Mem,
                                 offset: usize,
                                 size: usize,
                                 data: &mut [T]) -> Result<Event, Error> {
        let mut event_id: ll::Event = ptr::null_mut();
        let host_ptr = &mut data[0] as *mut T;
        let wait_list: Vec<ll::Event> = buf.read_events();
        let wait_size = wait_list.len() as u32;

        let wait_list_ptr: *const ll::Event = match wait_size {
            0 => ptr::null(),
            _ => &wait_list[0],
        };

        unsafe {
            try!(Error::check(ll::clEnqueueReadBuffer(self.id,
                                                      buf.id,
                                                      0, // don't block,
                                                      offset as size_t,
                                                      size as size_t,
                                                      transmute(host_ptr),
                                                      wait_size,
                                                      wait_list_ptr,
                                                      &mut event_id)));
        }

        let event = Event{id: event_id};
        buf.register_read(event.clone());
        Ok(event)
    }

    pub fn create_buffer(self: &Self,
                               size_bytes: usize) -> Result<Mem, Error> {
        let id = try!(self.alloc_ll_buffer(0, size_bytes));
        let tr = Mem::new(id);
        Ok(tr)
    }

    pub fn create_buffer_from_slice<T: Sized>(self: &Self,
                                    slice: &[T]) -> Result<Mem, Error> {
        let num_bytes = slice.len() * size_of::<T>();
        let id = try!(self.alloc_ll_buffer_with_data(0, num_bytes, unsafe { transmute(&slice[0]) } ));
        let tr = Mem::new(id);
        Ok(tr)
    }

    pub fn write_buffer<'a, 'b, T: Sized + 'a>(self: &Self,
                                  mem: &'b mut Mem,
                                  slice: &'a [T]) -> Result<ReadLock<'a, T>, Error> {
        let size_bytes = size_of::<T>()*slice.len();
        let evt = try!(self.write_buffer_raw(mem,
                                             0,
                                             size_bytes,
                                             slice));
        Ok(ReadLock{ evt: evt, sl: slice })
    }

    pub fn read_buffer<'a, T: Sized>(self: &Self,
                                 mem: &Mem,
                                 slice: &'a mut [T]) -> Result<WriteLock<'a, T>, Error> {
        let size_bytes = size_of::<T>()*slice.len();
        let evt = try!(self.read_buffer_raw(mem,
                                            0,
                                            size_bytes,
                                            slice));
        Ok(WriteLock{ evt: evt, sl: slice })
    }

    fn next_mul(self: &Self, global: usize, local: usize) -> usize {
        if global % local == 0 {
            global
        } else {
            ((global/local)+1)*local
        }
    }

    pub fn run( self: &Self,
                kernel: &mut Kernel,
                local_size: (usize, usize, usize),
                global_size: (usize, usize, usize)) -> Result<Event, Error> {
        let events = kernel.run_events();
        let global_size_fixed = (self.next_mul(global_size.0, local_size.0),
                                self.next_mul(global_size.1, local_size.1),
                                self.next_mul(global_size.2, local_size.2));

        let local_vec = vec!(local_size.0 as size_t,
                             local_size.1 as size_t,
                             local_size.2 as size_t);
        let global_vec = vec!(global_size_fixed.0 as size_t,
                              global_size_fixed.1 as size_t,
                              global_size_fixed.2 as size_t);

        let mut event_id: ll::Event = ptr::null_mut();
        unsafe {
            try!(Error::check(ll::clEnqueueNDRangeKernel(self.id,
                                                         kernel.id,
                                                         3,
                                                         ptr::null(),
                                                         &global_vec[0],
                                                         &local_vec[0],
                                                         events.len() as u32,
                                                         &events[0],
                                                         &mut event_id)));
        }
        let event = Event{id: event_id};

        try!(kernel.register_run(event.clone()));
        Ok(event)
    }
}

impl Drop for CommandQueue {
    fn drop(self: &mut Self) -> () {
        unsafe {
            Error::check(ll::clReleaseCommandQueue(self.id))
                .ok().expect("Error releasing command queue");
        }
    }
}

impl Clone for CommandQueue {
    fn clone(self: &Self) -> Self {
        unsafe {
            Error::check(ll::clRetainCommandQueue(self.id))
                .ok().expect("Error retaining command queue");
        }
        CommandQueue{ id: self.id }
    }
}

<commit_msg>fixes crash when kernel.run_events() == 0<commit_after>extern crate libc;
use self::libc::*;

use low_level as ll;
use error::Error;
use device::Device;
use context::Context;
use event::{Event, EventStatus, EventLike};
use kernel::Kernel;
use mem::Mem;

use std::ptr;
use std::mem::{size_of, transmute};

pub struct CommandQueue {
    pub id: ll::CommandQueue,
}

pub struct ReadLock<'a, T: Sized + 'a> {
    #[allow(dead_code)]
    sl: &'a [T],
    evt: Event,
}

pub struct WriteLock<'a, T: Sized + 'a> {
    #[allow(dead_code)]
    sl: &'a mut [T],
    evt: Event,
}

impl<'a, T: Sized> Drop for ReadLock<'a, T> {
    fn drop(self: &mut Self) -> () {
        self.evt.wait().ok().expect(
            "error waiting to finish write for ReadLock");
    }
}

impl<'a, T: Sized> EventLike for ReadLock<'a, T> {
    fn status(self: &Self) -> Result<EventStatus, Error> {
        self.evt.status()
    }

    fn wait(self: &Self) -> Result<(), Error> {
        self.evt.wait()
    }
}

impl<'a, T: Sized> Drop for WriteLock<'a, T> {
    fn drop(self: &mut Self) -> () {
        self.evt.wait().ok().expect(
            "error waiting to finish write for ReadLock");
    }
}

impl<'a, T: Sized> EventLike for WriteLock<'a, T> {
    fn status(self: &Self) -> Result<EventStatus, Error> {
        self.evt.status()
    }

    fn wait(self: &Self) -> Result<(), Error> {
        self.evt.wait()
    }
}

impl CommandQueue {
    pub fn new(context: Context,
               device: Device) -> Result<CommandQueue, Error> {
        let props: c_ulong = if try!(device.out_of_order_supported()) {
            1
        } else {
            0
        };
        let id = unsafe {
            let mut err: i32 = 0;
            let id = ll::clCreateCommandQueue(context.id,
                                              device.id,
                                              props,
                                              &mut err);
            try!(Error::check(err));
            id
        };
        Ok(CommandQueue{ id: id })
    }

    fn get_info(self: &Self, param: u32) -> Result<*mut c_void, Error> {
        let mut tr: *mut c_void = ptr::null_mut();
        unsafe {
            try!(Error::check(ll::clGetCommandQueueInfo(self.id,
                                                   param,
                                                   size_of::<*mut c_void>() as size_t,
                                                   transmute(&mut tr),
                                                   ptr::null_mut())));
        }
        Ok(tr)
    }

    pub fn context(self: &Self) -> Result<Context, Error> {
        let val = try!(self.get_info(0x1090));
        let id = unsafe { transmute(val) };
        unsafe {
            try!(Error::check(ll::clRetainContext(id)));
        }
        Ok(Context{id: id})
    }

    pub fn device(self: &Self) -> Result<Device, Error> {
        let val = try!(self.get_info(0x1091));
        let id = unsafe { transmute(val) };
        Ok(Device{id: id})
    }

    fn alloc_ll_buffer(self: &Self,
                    flags: c_ulong,
                    size_bytes: usize) -> Result<ll::Mem, Error> {
        let context = try!(self.context());
        let id = unsafe {
            let mut err: i32 = 0;
            let id = ll::clCreateBuffer(context.id,
                                        flags,
                                        size_bytes as size_t,
                                        ptr::null_mut(),
                                        &mut err);
            try!(Error::check(err));
            id
        };
        Ok(id)
    }

    fn alloc_ll_buffer_with_data(self: &Self,
                                 flags: c_ulong,
                                 size_bytes: usize,
                                 data: *const c_void) -> Result<ll::Mem, Error> {
        let context = try!(self.context());
        let id = unsafe {
            let mut err: i32 = 0;
            let id = ll::clCreateBuffer(context.id,
                                        flags | (1 << 5),
                                        size_bytes as size_t,
                                        data,
                                        &mut err);
            try!(Error::check(err));
            id
        };
        Ok(id)
    }

    fn write_buffer_raw<T: Sized>(self: &Self,
                       buf: &mut Mem,
                       offset: usize,
                       size: usize,
                       data: &[T]) -> Result<Event, Error> {
        let mut event_id: ll::Event = ptr::null_mut();
        let host_ptr = &data[0] as *const T;
        let wait_list: Vec<ll::Event> = buf.write_events();
        let wait_size = wait_list.len() as u32;

        let wait_list_ptr: *const ll::Event = match wait_size {
            0 => ptr::null(),
            _ => &wait_list[0],
        };

        unsafe {
            try!(Error::check(ll::clEnqueueWriteBuffer(self.id,
                                                       buf.id,
                                                       0, // don't block
                                                       offset as size_t,
                                                       size as size_t,
                                                       transmute(host_ptr),
                                                       wait_size,
                                                       wait_list_ptr,
                                                       &mut event_id)));
        }

        let event = Event{id: event_id};
        buf.register_write(event.clone());
        Ok(event)
    }

    fn read_buffer_raw<T: Sized>(self: &Self,
                                 buf: &Mem,
                                 offset: usize,
                                 size: usize,
                                 data: &mut [T]) -> Result<Event, Error> {
        let mut event_id: ll::Event = ptr::null_mut();
        let host_ptr = &mut data[0] as *mut T;
        let wait_list: Vec<ll::Event> = buf.read_events();
        let wait_size = wait_list.len() as u32;

        let wait_list_ptr: *const ll::Event = match wait_size {
            0 => ptr::null(),
            _ => &wait_list[0],
        };

        unsafe {
            try!(Error::check(ll::clEnqueueReadBuffer(self.id,
                                                      buf.id,
                                                      0, // don't block,
                                                      offset as size_t,
                                                      size as size_t,
                                                      transmute(host_ptr),
                                                      wait_size,
                                                      wait_list_ptr,
                                                      &mut event_id)));
        }

        let event = Event{id: event_id};
        buf.register_read(event.clone());
        Ok(event)
    }

    pub fn create_buffer(self: &Self,
                               size_bytes: usize) -> Result<Mem, Error> {
        let id = try!(self.alloc_ll_buffer(0, size_bytes));
        let tr = Mem::new(id);
        Ok(tr)
    }

    pub fn create_buffer_from_slice<T: Sized>(self: &Self,
                                    slice: &[T]) -> Result<Mem, Error> {
        let num_bytes = slice.len() * size_of::<T>();
        let id = try!(self.alloc_ll_buffer_with_data(0, num_bytes, unsafe { transmute(&slice[0]) } ));
        let tr = Mem::new(id);
        Ok(tr)
    }

    pub fn write_buffer<'a, 'b, T: Sized + 'a>(self: &Self,
                                  mem: &'b mut Mem,
                                  slice: &'a [T]) -> Result<ReadLock<'a, T>, Error> {
        let size_bytes = size_of::<T>()*slice.len();
        let evt = try!(self.write_buffer_raw(mem,
                                             0,
                                             size_bytes,
                                             slice));
        Ok(ReadLock{ evt: evt, sl: slice })
    }

    pub fn read_buffer<'a, T: Sized>(self: &Self,
                                 mem: &Mem,
                                 slice: &'a mut [T]) -> Result<WriteLock<'a, T>, Error> {
        let size_bytes = size_of::<T>()*slice.len();
        let evt = try!(self.read_buffer_raw(mem,
                                            0,
                                            size_bytes,
                                            slice));
        Ok(WriteLock{ evt: evt, sl: slice })
    }

    fn next_mul(self: &Self, global: usize, local: usize) -> usize {
        if global % local == 0 {
            global
        } else {
            ((global/local)+1)*local
        }
    }

    pub fn run( self: &Self,
                kernel: &mut Kernel,
                local_size: (usize, usize, usize),
                global_size: (usize, usize, usize)) -> Result<Event, Error> {
        let events = kernel.run_events();
        let global_size_fixed = (self.next_mul(global_size.0, local_size.0),
                                self.next_mul(global_size.1, local_size.1),
                                self.next_mul(global_size.2, local_size.2));

        let local_vec = vec!(local_size.0 as size_t,
                             local_size.1 as size_t,
                             local_size.2 as size_t);
        let global_vec = vec!(global_size_fixed.0 as size_t,
                              global_size_fixed.1 as size_t,
                              global_size_fixed.2 as size_t);

        let mut event_id: ll::Event = ptr::null_mut();
        unsafe {
            let events_val = if events.len() > 0 {
                transmute(&events[0])
            } else {
                ptr::null_mut()
            };
            try!(Error::check(ll::clEnqueueNDRangeKernel(self.id,
                                                         kernel.id,
                                                         3,
                                                         ptr::null(),
                                                         &global_vec[0],
                                                         &local_vec[0],
                                                         events.len() as u32,
                                                         events_val,
                                                         &mut event_id)));
        }
        let event = Event{id: event_id};

        try!(kernel.register_run(event.clone()));
        Ok(event)
    }
}

impl Drop for CommandQueue {
    fn drop(self: &mut Self) -> () {
        unsafe {
            Error::check(ll::clReleaseCommandQueue(self.id))
                .ok().expect("Error releasing command queue");
        }
    }
}

impl Clone for CommandQueue {
    fn clone(self: &Self) -> Self {
        unsafe {
            Error::check(ll::clRetainCommandQueue(self.id))
                .ok().expect("Error retaining command queue");
        }
        CommandQueue{ id: self.id }
    }
}

<|endoftext|>


<commit_before>extern crate rosc;
extern crate byteorder;

use byteorder::{ByteOrder, BigEndian};
use std::mem;
use rosc::{types, errors, decoder, utils};

#[test]
fn test_decode_no_args() {
    let raw_addr = "/some/valid/address/4";
    let addr = pad(raw_addr.as_bytes());
    // args
    let type_tags = pad(b",");
    let merged: Vec<u8> = addr.into_iter()
        .chain(type_tags.into_iter())
        .collect();
    let osc_packet: Result<types::OscPacket, errors::OscError> = decoder::decode(&merged);
    assert!(osc_packet.is_ok());
    match osc_packet {
        Ok(types::OscPacket::Message(msg)) => {
                assert_eq!(raw_addr, msg.addr);
                assert!(msg.args.is_none());
        },
        Ok(_) => panic!("OscMessage was expected!"),
        Err(e) => panic!(e)
    }
}

#[test]
fn test_decode_args() {
    let raw_addr = "/another/valid/address/123";
    let addr = pad(raw_addr.as_bytes());
    // args
    let f = 3.1415f32;
    let mut f_bytes: [u8; 4] = [0u8; 4];
    BigEndian::write_f32(&mut f_bytes, f);
    assert_eq!(BigEndian::read_f32(&f_bytes), f);

    let d = 3.14159265359f64;
    let mut d_bytes: [u8; 8] = [0u8; 8];
    BigEndian::write_f64(&mut d_bytes, d);
    assert_eq!(BigEndian::read_f64(&d_bytes), d);

    let i = 12345678i32;
    let i_bytes: [u8; 4] = unsafe {mem::transmute(i.to_be())};

    let l = -1234567891011i64;
    let h_bytes: [u8; 8] = unsafe {mem::transmute(l.to_be())};

    let type_tags = pad(b",fdih");

    let args: Vec<u8> = f_bytes.iter()
        .chain(d_bytes.iter())
        .chain(i_bytes.iter())
        .chain(h_bytes.iter())
        .map(|x| *x)
        .collect::<Vec<u8>>();

    let merged: Vec<u8> = addr.into_iter()
        .chain(type_tags.into_iter())
        .chain(args)
        .collect::<Vec<u8>>();
    let osc_packet: Result<types::OscPacket, errors::OscError> = decoder::decode(&merged);
    assert!(osc_packet.is_ok());
    match osc_packet {
        Ok(packet) => match packet {
            types::OscPacket::Message(msg) => {
                assert_eq!(raw_addr, msg.addr);
                match msg.args {
                    Some(args) => {
                        for arg in args {
                            match arg {
                                types::OscType::OscInt(x) => assert_eq!(i,x),
                                types::OscType::OscLong(x) => assert_eq!(l,x),
                                types::OscType::OscFloat(x) => assert_eq!(f, x),
                                types::OscType::OscDouble(x) => assert_eq!(d, x),
                                _ => panic!()
                            }
                        }
                    },
                    None => panic!()
                }
            },
            _ => panic!()
        },
        Err(e) => panic!(e)
    }
}

fn pad(data: &[u8]) -> Vec<u8> {
    let pad_len: usize = utils::pad(data.len() as u64) as usize;
    let mut v: Vec<u8> = data.iter().cloned().collect();
    while v.len() < pad_len {
        v.push(0u8);
    }
    v
}<commit_msg>Describe test messages<commit_after>extern crate rosc;
extern crate byteorder;

use byteorder::{ByteOrder, BigEndian};
use std::mem;
use rosc::{types, errors, decoder, utils};

#[test]
fn test_decode_no_args() {
    // message to build: /some/valid/address/4 ,
    let raw_addr = "/some/valid/address/4";
    let addr = pad(raw_addr.as_bytes());
    // args
    let type_tags = pad(b",");
    let merged: Vec<u8> = addr.into_iter()
        .chain(type_tags.into_iter())
        .collect();
    let osc_packet: Result<types::OscPacket, errors::OscError> = decoder::decode(&merged);
    assert!(osc_packet.is_ok());
    match osc_packet {
        Ok(types::OscPacket::Message(msg)) => {
                assert_eq!(raw_addr, msg.addr);
                assert!(msg.args.is_none());
        },
        Ok(_) => panic!("OscMessage was expected!"),
        Err(e) => panic!(e)
    }
}

#[test]
fn test_decode_args() {
    // /another/valid/address/123 ,fdih 3.1415 3.14159265359 12345678i32
    // -1234567891011
    let raw_addr = "/another/valid/address/123";
    let addr = pad(raw_addr.as_bytes());
    // args
    let f = 3.1415f32;
    let mut f_bytes: [u8; 4] = [0u8; 4];
    BigEndian::write_f32(&mut f_bytes, f);
    assert_eq!(BigEndian::read_f32(&f_bytes), f);

    let d = 3.14159265359f64;
    let mut d_bytes: [u8; 8] = [0u8; 8];
    BigEndian::write_f64(&mut d_bytes, d);
    assert_eq!(BigEndian::read_f64(&d_bytes), d);

    let i = 12345678i32;
    let i_bytes: [u8; 4] = unsafe {mem::transmute(i.to_be())};

    let l = -1234567891011i64;
    let h_bytes: [u8; 8] = unsafe {mem::transmute(l.to_be())};

    let type_tags = pad(b",fdih");

    let args: Vec<u8> = f_bytes.iter()
        .chain(d_bytes.iter())
        .chain(i_bytes.iter())
        .chain(h_bytes.iter())
        .map(|x| *x)
        .collect::<Vec<u8>>();

    let merged: Vec<u8> = addr.into_iter()
        .chain(type_tags.into_iter())
        .chain(args)
        .collect::<Vec<u8>>();
    let osc_packet: Result<types::OscPacket, errors::OscError> = decoder::decode(&merged);
    assert!(osc_packet.is_ok());
    match osc_packet {
        Ok(packet) => match packet {
            types::OscPacket::Message(msg) => {
                assert_eq!(raw_addr, msg.addr);
                match msg.args {
                    Some(args) => {
                        for arg in args {
                            match arg {
                                types::OscType::OscInt(x) => assert_eq!(i,x),
                                types::OscType::OscLong(x) => assert_eq!(l,x),
                                types::OscType::OscFloat(x) => assert_eq!(f, x),
                                types::OscType::OscDouble(x) => assert_eq!(d, x),
                                _ => panic!()
                            }
                        }
                    },
                    None => panic!()
                }
            },
            _ => panic!()
        },
        Err(e) => panic!(e)
    }
}

fn pad(data: &[u8]) -> Vec<u8> {
    let pad_len: usize = utils::pad(data.len() as u64) as usize;
    let mut v: Vec<u8> = data.iter().cloned().collect();
    while v.len() < pad_len {
        v.push(0u8);
    }
    v
}<|endoftext|>


<commit_before>use header::{Header, HeaderFormat};
use std::fmt::{mod, Show};
use super::util::from_one_raw_str;
use std::from_str::FromStr;
use time::{Tm, strptime};

// Egh, replace as soon as something better than time::Tm exists.
/// The `Date` header field.
#[deriving(PartialEq, Clone)]
pub struct Date(pub Tm);

impl Header for Date {
    fn header_name(_: Option<Date>) -> &'static str {
        "Date"
    }

    fn parse_header(raw: &[Vec<u8>]) -> Option<Date> {
        from_one_raw_str(raw)
    }
}

impl HeaderFormat for Date {
    fn fmt_header(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        self.fmt(fmt)
    }
}

impl fmt::Show for Date {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let Date(ref tm) = *self;
        // bummer that tm.strftime allocates a string. It would nice if it
        // returned a Show instead, since I don't need the String here
        write!(fmt, "{}", tm.to_utc().rfc822())
    }
}

impl FromStr for Date {
    //    Prior to 1995, there were three different formats commonly used by
    //   servers to communicate timestamps.  For compatibility with old
    //   implementations, all three are defined here.  The preferred format is
    //   a fixed-length and single-zone subset of the date and time
    //   specification used by the Internet Message Format [RFC5322].
    //
    //     HTTP-date    = IMF-fixdate / obs-date
    //
    //   An example of the preferred format is
    //
    //     Sun, 06 Nov 1994 08:49:37 GMT    ; IMF-fixdate
    //
    //   Examples of the two obsolete formats are
    //
    //     Sunday, 06-Nov-94 08:49:37 GMT   ; obsolete RFC 850 format
    //     Sun Nov  6 08:49:37 1994         ; ANSI C's asctime() format
    //
    //   A recipient that parses a timestamp value in an HTTP header field
    //   MUST accept all three HTTP-date formats.  When a sender generates a
    //   header field that contains one or more timestamps defined as
    //   HTTP-date, the sender MUST generate those timestamps in the
    //   IMF-fixdate format.
    fn from_str(s: &str) -> Option<Date> {
        strptime(s, "%a, %d %b %Y %T %Z").or_else(|_| {
            strptime(s, "%A, %d-%b-%y %T %Z")
        }).or_else(|_| {
            strptime(s, "%c")
        }).ok().map(|tm| Date(tm))
    }
}

<commit_msg>dont convert Date to utc if already utc<commit_after>use header::{Header, HeaderFormat};
use std::fmt::{mod, Show};
use super::util::from_one_raw_str;
use std::from_str::FromStr;
use time::{Tm, strptime};

// Egh, replace as soon as something better than time::Tm exists.
/// The `Date` header field.
#[deriving(PartialEq, Clone)]
pub struct Date(pub Tm);

impl Header for Date {
    fn header_name(_: Option<Date>) -> &'static str {
        "Date"
    }

    fn parse_header(raw: &[Vec<u8>]) -> Option<Date> {
        from_one_raw_str(raw)
    }
}

impl HeaderFormat for Date {
    fn fmt_header(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        self.fmt(fmt)
    }
}

impl fmt::Show for Date {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let Date(ref tm) = *self;
        // bummer that tm.strftime allocates a string. It would nice if it
        // returned a Show instead, since I don't need the String here
        match tm.tm_gmtoff {
            0 => tm.rfc822().fmt(fmt),
            _ => tm.to_utc().rfc822().fmt(fmt)
        }
    }
}

impl FromStr for Date {
    //    Prior to 1995, there were three different formats commonly used by
    //   servers to communicate timestamps.  For compatibility with old
    //   implementations, all three are defined here.  The preferred format is
    //   a fixed-length and single-zone subset of the date and time
    //   specification used by the Internet Message Format [RFC5322].
    //
    //     HTTP-date    = IMF-fixdate / obs-date
    //
    //   An example of the preferred format is
    //
    //     Sun, 06 Nov 1994 08:49:37 GMT    ; IMF-fixdate
    //
    //   Examples of the two obsolete formats are
    //
    //     Sunday, 06-Nov-94 08:49:37 GMT   ; obsolete RFC 850 format
    //     Sun Nov  6 08:49:37 1994         ; ANSI C's asctime() format
    //
    //   A recipient that parses a timestamp value in an HTTP header field
    //   MUST accept all three HTTP-date formats.  When a sender generates a
    //   header field that contains one or more timestamps defined as
    //   HTTP-date, the sender MUST generate those timestamps in the
    //   IMF-fixdate format.
    fn from_str(s: &str) -> Option<Date> {
        strptime(s, "%a, %d %b %Y %T %Z").or_else(|_| {
            strptime(s, "%A, %d-%b-%y %T %Z")
        }).or_else(|_| {
            strptime(s, "%c")
        }).ok().map(|tm| Date(tm))
    }
}

<|endoftext|>


<commit_before>use super::{SeasonState, SeriesConfig};
use backend::{AnimeEntry, AnimeInfo, SyncBackend};
use error::SeriesError;
use input::{self, Answer};
use regex::Regex;
use std::borrow::Cow;
use std::collections::HashMap;
use std::fs;
use std::ops::Range;
use std::path::{Path, PathBuf};
use toml;

pub struct FolderData {
    pub series_info: SeriesInfo,
    pub savefile: SaveData,
    pub subseries: String,
    pub path: PathBuf,
}

impl FolderData {
    const DEFAULT_SUBSERIES_NAME: &'static str = "main";

    pub fn load_dir(path: &Path, subseries: Option<String>) -> Result<FolderData, SeriesError> {
        let subseries = subseries.unwrap_or_else(|| FolderData::DEFAULT_SUBSERIES_NAME.into());

        let mut savefile = SaveData::from_dir(path)?;
        let series_info = FolderData::load_series_info(path, &mut savefile, &subseries)?;

        Ok(FolderData {
            series_info,
            savefile,
            subseries,
            path: PathBuf::from(path),
        })
    }

    fn load_series_info(
        path: &Path,
        savefile: &mut SaveData,
        subseries: &str,
    ) -> Result<SeriesInfo, SeriesError> {
        let mut ep_data = parse_episode_files_until_valid(path, &mut savefile.episode_matcher)?;
        let subseries_data = savefile.get_mut_subseries_entry(subseries);

        if let Some(info) = SeriesInfo::select_from_subseries(&mut ep_data, &subseries_data) {
            return Ok(info);
        }

        let info = prompt_select_series_info(ep_data)?;
        subseries_data.files_title = Some(info.name.clone());

        Ok(info)
    }

    pub fn populate_season_data<B>(&mut self, config: &SeriesConfig<B>) -> Result<(), SeriesError>
    where
        B: SyncBackend,
    {
        let num_seasons = self.seasons().len();

        if num_seasons > config.season_num {
            return Ok(());
        }

        for cur_season in num_seasons..=config.season_num {
            let info = AnimeInfo::default();
            let entry = AnimeEntry::new(info);

            let mut season = SeasonState {
                state: entry,
                needs_info: true,
                needs_sync: config.offline_mode,
            };

            season.sync_info_from_remote(config, &self, cur_season)?;

            let subseries = self
                .savefile
                .get_mut_subseries_entry(self.subseries.clone());

            subseries.season_states.push(season);
        }

        Ok(())
    }

    pub fn calculate_season_offset(&self, mut range: Range<usize>) -> u32 {
        let num_seasons = self.seasons().len();
        range.start = num_seasons.min(range.start);
        range.end = num_seasons.min(range.end);

        let mut offset = 0;
        let seasons = self.seasons();

        for i in range {
            let season = &seasons[i];

            match season.state.info.episodes {
                Some(eps) => offset += eps,
                None => return offset,
            }
        }

        offset
    }

    pub fn try_remove_dir(&self) {
        let path = self.path.to_string_lossy();

        println!("WARNING: {} will be deleted", path);
        println!("is this ok? (y/N)");

        match input::read_yn(Answer::No) {
            Ok(true) => match fs::remove_dir_all(&self.path) {
                Ok(_) => (),
                Err(err) => {
                    eprintln!("failed to remove directory: {}", err);
                }
            },
            Ok(false) => (),
            Err(err) => {
                eprintln!("failed to read input: {}", err);
            }
        }
    }

    pub fn get_episode(&self, episode: u32) -> Result<&PathBuf, SeriesError> {
        self.series_info.get_episode(episode)
    }

    pub fn save(&self) -> Result<(), SeriesError> {
        self.savefile.write_to_file()
    }

    pub fn seasons(&self) -> &Vec<SeasonState> {
        &self.savefile.subseries[&self.subseries].season_states
    }

    pub fn seasons_mut(&mut self) -> &mut Vec<SeasonState> {
        &mut self
            .savefile
            .get_mut_subseries_entry(self.subseries.clone())
            .season_states
    }
}

pub type SubSeriesName = String;

#[derive(Serialize, Deserialize)]
pub struct SaveData {
    pub episode_matcher: Option<String>,
    #[serde(flatten)]
    pub subseries: HashMap<SubSeriesName, SubSeriesData>,
    #[serde(skip)]
    pub path: PathBuf,
}

impl SaveData {
    const DATA_FILE_NAME: &'static str = ".anup";

    pub fn new(path: PathBuf) -> SaveData {
        SaveData {
            episode_matcher: None,
            subseries: HashMap::new(),
            path,
        }
    }

    pub fn get_mut_subseries_entry<S>(&mut self, name: S) -> &mut SubSeriesData
    where
        S: Into<String>,
    {
        self.subseries
            .entry(name.into())
            .or_insert_with(SubSeriesData::new)
    }

    pub fn from_dir(path: &Path) -> Result<SaveData, SeriesError> {
        let path = PathBuf::from(path).join(SaveData::DATA_FILE_NAME);

        if !path.exists() {
            return Ok(SaveData::new(path));
        }

        let file_contents = fs::read_to_string(&path)?;

        let mut save_data: SaveData = toml::from_str(&file_contents)?;
        save_data.path = path;

        Ok(save_data)
    }

    pub fn write_to_file(&self) -> Result<(), SeriesError> {
        let toml = toml::to_string(self)?;
        fs::write(&self.path, toml)?;

        Ok(())
    }
}

#[derive(Serialize, Deserialize)]
pub struct SubSeriesData {
    pub files_title: Option<String>,
    pub season_states: Vec<SeasonState>,
}

impl SubSeriesData {
    pub fn new() -> SubSeriesData {
        SubSeriesData {
            files_title: None,
            season_states: Vec::new(),
        }
    }
}

pub struct SeriesInfo {
    pub name: String,
    pub episodes: HashMap<u32, PathBuf>,
}

impl SeriesInfo {
    pub fn get_episode(&self, episode: u32) -> Result<&PathBuf, SeriesError> {
        self.episodes
            .get(&episode)
            .ok_or_else(|| SeriesError::EpisodeNotFound(episode))
    }

    pub fn select_from_subseries(
        ep_data: &mut SeriesEpisodes,
        subseries: &SubSeriesData,
    ) -> Option<SeriesInfo> {
        if let Some(name) = &subseries.files_title {
            let entry = ep_data.remove_entry(name);

            if let Some((name, episodes)) = entry {
                return Some(SeriesInfo { name, episodes });
            }
        }

        None
    }
}

struct EpisodeFile {
    series_name: String,
    episode_num: u32,
}

impl EpisodeFile {
    fn parse(path: &Path, matcher: &Regex) -> Result<EpisodeFile, SeriesError> {
        // Replace certain characters with spaces since they can prevent proper series
        // identification or prevent it from being found on a sync backend
        let filename = path
            .file_name()
            .and_then(|path| path.to_str())
            .ok_or(SeriesError::UnableToGetFilename)?;

        let caps = matcher
            .captures(&filename)
            .ok_or(SeriesError::EpisodeRegexCaptureFailed)?;

        let series_name = caps
            .name("name")
            .ok_or_else(|| SeriesError::UnknownRegexCapture("name".into()))
            .map(|name| EpisodeFile::cleanup_title(name.as_str()))?;

        let episode = match caps.name("episode") {
            Some(capture) => capture
                .as_str()
                .parse()
                .map_err(SeriesError::EpisodeNumParseFailed)?,
            None => 1,
        };

        Ok(EpisodeFile {
            series_name,
            episode_num: episode,
        })
    }

    fn cleanup_title(title: &str) -> String {
        let mut clean_title = String::with_capacity(title.len());
        let mut last_char = ' ';

        for ch in title.chars() {
            match ch {
                '.' => clean_title.push(' '),
                '_' => clean_title.push(' '),
                ' ' if last_char == ' ' => (),
                ch => clean_title.push(ch),
            }

            last_char = ch;
        }

        clean_title
    }
}

pub type EpisodePaths = HashMap<u32, PathBuf>;
pub type SeriesEpisodes = HashMap<String, EpisodePaths>;

pub fn prompt_select_series_info(info: SeriesEpisodes) -> Result<SeriesInfo, SeriesError> {
    if info.is_empty() {
        return Err(SeriesError::NoSeriesFound);
    }

    let mut info = info
        .into_iter()
        .map(|(name, eps)| SeriesInfo {
            name,
            episodes: eps,
        })
        .collect::<Vec<_>>();

    if info.len() == 1 {
        return Ok(info.swap_remove(0));
    }

    println!("multiple series found in directory");
    println!("please enter the number next to the episode files you want to use:");

    for (i, series) in info.iter().enumerate() {
        println!("{} [{}]", 1 + i, series.name);
    }

    let index = input::read_range(1, info.len())? - 1;
    let series = info.swap_remove(index);

    Ok(series)
}

// This default pattern will match episodes in several common formats, such as:
// [Group] Series Name - 01.mkv
// [Group]_Series_Name_-_01.mkv
// [Group].Series.Name.-.01.mkv
// [Group] Series Name - 01 [tag 1][tag 2].mkv
// [Group]_Series_Name_-_01_[tag1][tag2].mkv
// [Group].Series.Name.-.01.[tag1][tag2].mkv
// Series Name - 01.mkv
// Series_Name_-_01.mkv
// Series.Name.-.01.mkv
const EP_FORMAT_REGEX: &str =
    r"(?:\[.+?\](?:_+|\.+|\s*))?(?P<name>.+?)(?:\s*|_*|\.*)-(?:\s*|_*|\.*)(?P<episode>\d+).*?\..+?";

fn format_episode_parser_regex<'a, S>(pattern: Option<S>) -> Result<Cow<'a, Regex>, SeriesError>
where
    S: AsRef<str>,
{
    lazy_static! {
        static ref EP_FORMAT: Regex = Regex::new(EP_FORMAT_REGEX).unwrap();
    }

    match pattern {
        Some(pattern) => {
            let pattern = pattern
                .as_ref()
                .replace("{name}", "(?P<name>.+?)")
                .replace("{episode}", r"(?P<episode>\d+)");

            let regex = Regex::new(&pattern)?;
            Ok(Cow::Owned(regex))
        }
        None => Ok(Cow::Borrowed(&*EP_FORMAT)),
    }
}

pub fn parse_episode_files<S>(
    path: &Path,
    pattern: Option<S>,
) -> Result<SeriesEpisodes, SeriesError>
where
    S: AsRef<str>,
{
    if !path.is_dir() {
        return Err(SeriesError::NotADirectory(path.to_string_lossy().into()));
    }

    let pattern = format_episode_parser_regex(pattern)?;
    let mut data = HashMap::with_capacity(1);

    for entry in fs::read_dir(path).map_err(SeriesError::Io)? {
        let entry = entry.map_err(SeriesError::Io)?.path();

        if !entry.is_file() {
            continue;
        }

        let episode = match EpisodeFile::parse(&entry, pattern.as_ref()) {
            Ok(episode) => episode,
            Err(SeriesError::EpisodeRegexCaptureFailed) => continue,
            Err(err) => return Err(err),
        };

        let series = data
            .entry(episode.series_name)
            .or_insert_with(|| HashMap::with_capacity(1));

        series.insert(episode.episode_num, entry);
    }

    if data.is_empty() {
        return Err(SeriesError::NoSeriesFound);
    }

    Ok(data)
}

pub fn parse_episode_files_until_valid<S>(
    path: &Path,
    pattern: &mut Option<S>,
) -> Result<SeriesEpisodes, SeriesError>
where
    S: AsRef<str> + From<String>,
{
    loop {
        match parse_episode_files(path, pattern.as_ref()) {
            Ok(data) => break Ok(data),
            Err(SeriesError::NoSeriesFound) => {
                eprintln!("error: no series found");
                println!("you will now be prompted to enter a custom regex pattern");
                println!("when entering the pattern, please mark the series name and episode number with {{name}} and {{episode}}, respectively");
                println!("note: if you're trying to play a single file (like a movie) and it doesn't have an episode number, you can omit the {{episode}} marker");
                println!("example:");
                println!("  filename: [SubGroup] Series Name - Ep01 [1080p].mkv");
                println!(r"  pattern: \[.+?\] {{name}} - Ep{{episode}} \[1080p\].mkv");
                println!("please enter your custom pattern:");

                *pattern = Some(input::read_line()?.into());
            }
            Err(err @ SeriesError::Regex(_)) | Err(err @ SeriesError::UnknownRegexCapture(_)) => {
                eprintln!("error parsing regex pattern: {}", err);
                println!("please try again:");

                *pattern = Some(input::read_line()?.into());
            }
            Err(err) => return Err(err),
        }
    }
}
<commit_msg>[series] Further optimize series episode parsing<commit_after>use super::{SeasonState, SeriesConfig};
use backend::{AnimeEntry, AnimeInfo, SyncBackend};
use error::SeriesError;
use input::{self, Answer};
use regex::Regex;
use std::borrow::Cow;
use std::collections::HashMap;
use std::fs;
use std::ops::Range;
use std::path::{Path, PathBuf};
use toml;

pub struct FolderData {
    pub series_info: SeriesInfo,
    pub savefile: SaveData,
    pub subseries: String,
    pub path: PathBuf,
}

impl FolderData {
    const DEFAULT_SUBSERIES_NAME: &'static str = "main";

    pub fn load_dir(path: &Path, subseries: Option<String>) -> Result<FolderData, SeriesError> {
        let subseries = subseries.unwrap_or_else(|| FolderData::DEFAULT_SUBSERIES_NAME.into());

        let mut savefile = SaveData::from_dir(path)?;
        let series_info = FolderData::load_series_info(path, &mut savefile, &subseries)?;

        Ok(FolderData {
            series_info,
            savefile,
            subseries,
            path: PathBuf::from(path),
        })
    }

    fn load_series_info(
        path: &Path,
        savefile: &mut SaveData,
        subseries: &str,
    ) -> Result<SeriesInfo, SeriesError> {
        let mut ep_data = parse_episode_files_until_valid(path, &mut savefile.episode_matcher)?;
        let subseries_data = savefile.get_mut_subseries_entry(subseries);

        if let Some(info) = SeriesInfo::select_from_subseries(&mut ep_data, &subseries_data) {
            return Ok(info);
        }

        let info = prompt_select_series_info(ep_data)?;
        subseries_data.files_title = Some(info.name.clone());

        Ok(info)
    }

    pub fn populate_season_data<B>(&mut self, config: &SeriesConfig<B>) -> Result<(), SeriesError>
    where
        B: SyncBackend,
    {
        let num_seasons = self.seasons().len();

        if num_seasons > config.season_num {
            return Ok(());
        }

        for cur_season in num_seasons..=config.season_num {
            let info = AnimeInfo::default();
            let entry = AnimeEntry::new(info);

            let mut season = SeasonState {
                state: entry,
                needs_info: true,
                needs_sync: config.offline_mode,
            };

            season.sync_info_from_remote(config, &self, cur_season)?;

            let subseries = self
                .savefile
                .get_mut_subseries_entry(self.subseries.clone());

            subseries.season_states.push(season);
        }

        Ok(())
    }

    pub fn calculate_season_offset(&self, mut range: Range<usize>) -> u32 {
        let num_seasons = self.seasons().len();
        range.start = num_seasons.min(range.start);
        range.end = num_seasons.min(range.end);

        let mut offset = 0;
        let seasons = self.seasons();

        for i in range {
            let season = &seasons[i];

            match season.state.info.episodes {
                Some(eps) => offset += eps,
                None => return offset,
            }
        }

        offset
    }

    pub fn try_remove_dir(&self) {
        let path = self.path.to_string_lossy();

        println!("WARNING: {} will be deleted", path);
        println!("is this ok? (y/N)");

        match input::read_yn(Answer::No) {
            Ok(true) => match fs::remove_dir_all(&self.path) {
                Ok(_) => (),
                Err(err) => {
                    eprintln!("failed to remove directory: {}", err);
                }
            },
            Ok(false) => (),
            Err(err) => {
                eprintln!("failed to read input: {}", err);
            }
        }
    }

    pub fn get_episode(&self, episode: u32) -> Result<&PathBuf, SeriesError> {
        self.series_info.get_episode(episode)
    }

    pub fn save(&self) -> Result<(), SeriesError> {
        self.savefile.write_to_file()
    }

    pub fn seasons(&self) -> &Vec<SeasonState> {
        &self.savefile.subseries[&self.subseries].season_states
    }

    pub fn seasons_mut(&mut self) -> &mut Vec<SeasonState> {
        &mut self
            .savefile
            .get_mut_subseries_entry(self.subseries.clone())
            .season_states
    }
}

pub type SubSeriesName = String;

#[derive(Serialize, Deserialize)]
pub struct SaveData {
    pub episode_matcher: Option<String>,
    #[serde(flatten)]
    pub subseries: HashMap<SubSeriesName, SubSeriesData>,
    #[serde(skip)]
    pub path: PathBuf,
}

impl SaveData {
    const DATA_FILE_NAME: &'static str = ".anup";

    pub fn new(path: PathBuf) -> SaveData {
        SaveData {
            episode_matcher: None,
            subseries: HashMap::new(),
            path,
        }
    }

    pub fn get_mut_subseries_entry<S>(&mut self, name: S) -> &mut SubSeriesData
    where
        S: Into<String>,
    {
        self.subseries
            .entry(name.into())
            .or_insert_with(SubSeriesData::new)
    }

    pub fn from_dir(path: &Path) -> Result<SaveData, SeriesError> {
        let path = PathBuf::from(path).join(SaveData::DATA_FILE_NAME);

        if !path.exists() {
            return Ok(SaveData::new(path));
        }

        let file_contents = fs::read_to_string(&path)?;

        let mut save_data: SaveData = toml::from_str(&file_contents)?;
        save_data.path = path;

        Ok(save_data)
    }

    pub fn write_to_file(&self) -> Result<(), SeriesError> {
        let toml = toml::to_string(self)?;
        fs::write(&self.path, toml)?;

        Ok(())
    }
}

#[derive(Serialize, Deserialize)]
pub struct SubSeriesData {
    pub files_title: Option<String>,
    pub season_states: Vec<SeasonState>,
}

impl SubSeriesData {
    pub fn new() -> SubSeriesData {
        SubSeriesData {
            files_title: None,
            season_states: Vec::new(),
        }
    }
}

pub struct SeriesInfo {
    pub name: String,
    pub episodes: HashMap<u32, PathBuf>,
}

impl SeriesInfo {
    pub fn get_episode(&self, episode: u32) -> Result<&PathBuf, SeriesError> {
        self.episodes
            .get(&episode)
            .ok_or_else(|| SeriesError::EpisodeNotFound(episode))
    }

    pub fn select_from_subseries(
        ep_data: &mut SeriesEpisodes,
        subseries: &SubSeriesData,
    ) -> Option<SeriesInfo> {
        if let Some(name) = &subseries.files_title {
            let entry = ep_data.remove_entry(name);

            if let Some((name, episodes)) = entry {
                return Some(SeriesInfo { name, episodes });
            }
        }

        None
    }
}

struct EpisodeFile {
    series_name: String,
    episode_num: u32,
}

impl EpisodeFile {
    fn parse(path: &Path, matcher: &Regex) -> Result<EpisodeFile, SeriesError> {
        // Replace certain characters with spaces since they can prevent proper series
        // identification or prevent it from being found on a sync backend
        let filename = path
            .file_name()
            .and_then(|path| path.to_str())
            .ok_or(SeriesError::UnableToGetFilename)?;

        let caps = matcher
            .captures(&filename)
            .ok_or(SeriesError::EpisodeRegexCaptureFailed)?;

        let series_name = caps
            .name("name")
            .ok_or_else(|| SeriesError::UnknownRegexCapture("name".into()))
            .map(|name| EpisodeFile::cleanup_title(name.as_str()))?;

        let episode = match caps.name("episode") {
            Some(capture) => capture
                .as_str()
                .parse()
                .map_err(SeriesError::EpisodeNumParseFailed)?,
            None => 1,
        };

        Ok(EpisodeFile {
            series_name,
            episode_num: episode,
        })
    }

    fn cleanup_title(title: &str) -> String {
        let mut clean_title = String::from(title);
        let bytes = unsafe { clean_title.as_bytes_mut() };

        for byte in bytes {
            match byte {
                b'.' | b'_' => *byte = b' ',
                _ => (),
            }
        }

        clean_title
    }
}

pub type EpisodePaths = HashMap<u32, PathBuf>;
pub type SeriesEpisodes = HashMap<String, EpisodePaths>;

pub fn prompt_select_series_info(info: SeriesEpisodes) -> Result<SeriesInfo, SeriesError> {
    if info.is_empty() {
        return Err(SeriesError::NoSeriesFound);
    }

    let mut info = info
        .into_iter()
        .map(|(name, eps)| SeriesInfo {
            name,
            episodes: eps,
        })
        .collect::<Vec<_>>();

    if info.len() == 1 {
        return Ok(info.swap_remove(0));
    }

    println!("multiple series found in directory");
    println!("please enter the number next to the episode files you want to use:");

    for (i, series) in info.iter().enumerate() {
        println!("{} [{}]", 1 + i, series.name);
    }

    let index = input::read_range(1, info.len())? - 1;
    let series = info.swap_remove(index);

    Ok(series)
}

// This default pattern will match episodes in several common formats, such as:
// [Group] Series Name - 01.mkv
// [Group]_Series_Name_-_01.mkv
// [Group].Series.Name.-.01.mkv
// [Group] Series Name - 01 [tag 1][tag 2].mkv
// [Group]_Series_Name_-_01_[tag1][tag2].mkv
// [Group].Series.Name.-.01.[tag1][tag2].mkv
// Series Name - 01.mkv
// Series_Name_-_01.mkv
// Series.Name.-.01.mkv
const EP_FORMAT_REGEX: &str =
    r"(?:\[.+?\](?:_+|\.+|\s*))?(?P<name>.+?)(?:\s*|_*|\.*)-(?:\s*|_*|\.*)(?P<episode>\d+).*?\..+?";

fn format_episode_parser_regex<'a, S>(pattern: Option<S>) -> Result<Cow<'a, Regex>, SeriesError>
where
    S: AsRef<str>,
{
    lazy_static! {
        static ref EP_FORMAT: Regex = Regex::new(EP_FORMAT_REGEX).unwrap();
    }

    match pattern {
        Some(pattern) => {
            let pattern = pattern
                .as_ref()
                .replace("{name}", "(?P<name>.+?)")
                .replace("{episode}", r"(?P<episode>\d+)");

            let regex = Regex::new(&pattern)?;
            Ok(Cow::Owned(regex))
        }
        None => Ok(Cow::Borrowed(&*EP_FORMAT)),
    }
}

pub fn parse_episode_files<S>(
    path: &Path,
    pattern: Option<S>,
) -> Result<SeriesEpisodes, SeriesError>
where
    S: AsRef<str>,
{
    if !path.is_dir() {
        return Err(SeriesError::NotADirectory(path.to_string_lossy().into()));
    }

    let pattern = format_episode_parser_regex(pattern)?;
    let mut data = HashMap::with_capacity(1);

    for entry in fs::read_dir(path).map_err(SeriesError::Io)? {
        let entry = entry.map_err(SeriesError::Io)?.path();

        if !entry.is_file() {
            continue;
        }

        let episode = match EpisodeFile::parse(&entry, pattern.as_ref()) {
            Ok(episode) => episode,
            Err(SeriesError::EpisodeRegexCaptureFailed) => continue,
            Err(err) => return Err(err),
        };

        let series = data
            .entry(episode.series_name)
            .or_insert_with(|| HashMap::with_capacity(1));

        series.insert(episode.episode_num, entry);
    }

    if data.is_empty() {
        return Err(SeriesError::NoSeriesFound);
    }

    Ok(data)
}

pub fn parse_episode_files_until_valid<S>(
    path: &Path,
    pattern: &mut Option<S>,
) -> Result<SeriesEpisodes, SeriesError>
where
    S: AsRef<str> + From<String>,
{
    loop {
        match parse_episode_files(path, pattern.as_ref()) {
            Ok(data) => break Ok(data),
            Err(SeriesError::NoSeriesFound) => {
                eprintln!("error: no series found");
                println!("you will now be prompted to enter a custom regex pattern");
                println!("when entering the pattern, please mark the series name and episode number with {{name}} and {{episode}}, respectively");
                println!("note: if you're trying to play a single file (like a movie) and it doesn't have an episode number, you can omit the {{episode}} marker");
                println!("example:");
                println!("  filename: [SubGroup] Series Name - Ep01 [1080p].mkv");
                println!(r"  pattern: \[.+?\] {{name}} - Ep{{episode}} \[1080p\].mkv");
                println!("please enter your custom pattern:");

                *pattern = Some(input::read_line()?.into());
            }
            Err(err @ SeriesError::Regex(_)) | Err(err @ SeriesError::UnknownRegexCapture(_)) => {
                eprintln!("error parsing regex pattern: {}", err);
                println!("please try again:");

                *pattern = Some(input::read_line()?.into());
            }
            Err(err) => return Err(err),
        }
    }
}
<|endoftext|>


<commit_before>use alloc::{Allocator, SliceWrapper, SliceWrapperMut};
use core::marker::PhantomData;
use super::BrotliAlloc;
use super::backward_references::BrotliEncoderParams;
pub trait Joinable<T:Send+'static>:Sized {
  fn join(self) -> Result<T, ()>;
}

pub struct CompressionThreadResult<Alloc:BrotliAlloc+Send+'static> where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  compressed: <Alloc as Allocator<u8>>::AllocatedMemory,
  compressed_size: usize,
  alloc: Alloc,
}
enum InternalSendAlloc<T:Send+'static, Alloc:BrotliAlloc+Send+'static, Join: Joinable<T>>
  where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  A(Alloc),
  Join(Join),
  Spawning(PhantomData<T>),
}
pub struct SendAlloc<T:Send+'static,
                     Alloc:BrotliAlloc +Send+'static,
                     Join:Joinable<T>>(InternalSendAlloc<T, Alloc, Join>)
  where <Alloc as Allocator<u8>>::AllocatedMemory: Send;

impl<T:Send+'static, Alloc:BrotliAlloc+Send+'static,Join:Joinable<T>> SendAlloc<T, Alloc, Join>
  where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  pub fn new(alloc: Alloc) -> Self {
    SendAlloc::<T, Alloc, Join>(InternalSendAlloc::A(alloc))
  }
  pub fn unwrap_or(self, other: Alloc) -> Alloc {
    match self.0 {
      InternalSendAlloc::A(alloc) => {
        alloc
      },
      InternalSendAlloc::Spawning(_) | InternalSendAlloc::Join(_) => {
        other
      },
    }
  }
  pub fn unwrap(self) -> Alloc {
    match self.0 {
      InternalSendAlloc::A(alloc) => {
        alloc
      },
      InternalSendAlloc::Join(_) | InternalSendAlloc::Spawning(_) => panic!("Item permanently borrowed/leaked"),
    }
  }
}

enum InternalOwned<T> {
  Item(T),
  Borrowed,
}

pub struct Owned<T>(InternalOwned<T>);
impl<T> Owned<T> {
  pub fn new(data:T) -> Self {
    Owned::<T>(InternalOwned::Item(data))
  }
  pub fn unwrap_or(self, other: T) -> T {
    if let InternalOwned::Item(x) = self.0 {
      x
    } else {
      other
    }
  }
  pub fn unwrap(self) -> T {
    if let InternalOwned::Item(x) = self.0 {
      x
    } else {
      panic!("Item permanently borrowed")
    }
  }
}

pub trait OwnedRetriever<SliceW: SliceWrapper<u8>+Send+'static> {
  fn view(&self) -> &[u8];
  fn unwrap(self) -> SliceW;
}

pub trait BatchSpawnable<T:Send+'static, Alloc:BrotliAlloc+Send+'static, SliceW:SliceWrapper<u8>+Send+'static>
  where <Alloc as Allocator<u8>>::AllocatedMemory:Send+'static
{
  type JoinHandle: Joinable<T>;
  type FinalJoinHandle: OwnedRetriever<SliceW>;
  // this function takes in a SendAlloc per thread and converts them all into JoinHandle
  // the input is borrowed until the joins complete
  // owned is set to borrowed
  // the final join handle is a r/w lock which will return the SliceW to the owner
  // the FinalJoinHandle is only to be called when each individual JoinHandle has been examined
  fn batch_spawn<F: FnOnce(usize, &SliceW, Alloc) -> T>(
    &mut self,
    input: &mut Owned<SliceW>,
    alloc_per_thread:&mut [SendAlloc<T, Alloc, Self::JoinHandle>],
    f: F,
  ) -> Self::FinalJoinHandle;
}

pub fn CompressMultiSlice<Alloc:BrotliAlloc+Send+'static,
                          Spawner:BatchSpawnable<CompressionThreadResult<Alloc>,
                                                 Alloc,
                                                 <Alloc as Allocator<u8>>::AllocatedMemory>> (
  params:&BrotliEncoderParams,
  input_slice: &[u8],
  output: &mut [u8],
  alloc_per_thread:&mut [SendAlloc<CompressionThreadResult<Alloc>, Alloc, Spawner::JoinHandle>],
  mut thread_spawner: Spawner,
) -> Result<usize, ()> where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  let input = if let InternalSendAlloc::A(ref mut alloc) = alloc_per_thread[0].0 {
    let mut input = <Alloc as Allocator<u8>>::alloc_cell(alloc, input_slice.len());
    input.slice_mut().clone_from_slice(input_slice);
    input
  } else {
    <Alloc as Allocator<u8>>::AllocatedMemory::default()
  };
  let mut owned_input = Owned::new(input);
  let ret = CompressMulti(params, &mut owned_input, output, alloc_per_thread,thread_spawner);
  if let InternalSendAlloc::A(ref mut alloc) = alloc_per_thread[0].0 {
    <Alloc as Allocator<u8>>::free_cell(alloc, owned_input.unwrap());
  }
  ret
}

pub fn CompressMulti<Alloc:BrotliAlloc+Send+'static,
                     SliceW: SliceWrapper<u8>+Send+'static,
                     Spawner:BatchSpawnable<CompressionThreadResult<Alloc>,
                                            Alloc,
                                            SliceW>> (
  params:&BrotliEncoderParams,
  owned_input: &mut Owned<SliceW>,
  output: &mut [u8],
  alloc_per_thread:&mut [SendAlloc<CompressionThreadResult<Alloc>, Alloc, Spawner::JoinHandle>],
  mut thread_spawner: Spawner,
) -> Result<usize, ()> where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  let (mut alloc, alloc_rest) = alloc_per_thread.split_at_mut(1);
  let mut retrieve_owned_input = thread_spawner.batch_spawn(owned_input, alloc_rest, |_index,_input,mut alloc|{
    let mem = <Alloc as Allocator<u8>>::alloc_cell(&mut alloc, 0);
    CompressionThreadResult::<Alloc>{
      compressed:mem,
      compressed_size:0,
      alloc:alloc,
    }
  });
  {
    let input = retrieve_owned_input.view();

  }
  Err(())
}

<commit_msg>finish first thread compression logic and hook it into thread spawning<commit_after>use alloc::{Allocator, SliceWrapper, SliceWrapperMut};
use core::marker::PhantomData;
use core::mem;
use super::BrotliAlloc;
use super::encode::{
    BrotliEncoderOperation,
    BrotliEncoderCreateInstance,
    BrotliEncoderDestroyInstance,
    BrotliEncoderCompressStream,
};
use core::ops::Range;
use super::backward_references::BrotliEncoderParams;
pub trait Joinable<T:Send+'static, U:Send+'static>:Sized {
  fn join(self) -> Result<T, U>;
}

pub struct CompressionThreadResult<Alloc:BrotliAlloc+Send+'static> where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  compressed: <Alloc as Allocator<u8>>::AllocatedMemory,
  compressed_size: usize,
  alloc: Alloc,
}
enum InternalSendAlloc<T:Send+'static, Alloc:BrotliAlloc+Send+'static, Join: Joinable<T, Alloc>>
  where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  A(Alloc),
  Join(Join),
  SpawningOrJoining(PhantomData<T>),
}
pub struct SendAlloc<T:Send+'static,
                     Alloc:BrotliAlloc +Send+'static,
                     Join:Joinable<T, Alloc>>(InternalSendAlloc<T, Alloc, Join>)
  where <Alloc as Allocator<u8>>::AllocatedMemory: Send;

impl<T:Send+'static, Alloc:BrotliAlloc+Send+'static,Join:Joinable<T, Alloc>> SendAlloc<T, Alloc, Join>
  where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  pub fn new(alloc: Alloc) -> Self {
    SendAlloc::<T, Alloc, Join>(InternalSendAlloc::A(alloc))
  }
  pub fn unwrap_or(self, other: Alloc) -> Alloc {
    match self.0 {
      InternalSendAlloc::A(alloc) => {
        alloc
      },
      InternalSendAlloc::SpawningOrJoining(_) | InternalSendAlloc::Join(_) => {
        other
      },
    }
  }
  pub fn unwrap(self) -> Alloc {
    match self.0 {
      InternalSendAlloc::A(alloc) => {
        alloc
      },
      InternalSendAlloc::SpawningOrJoining(_) | InternalSendAlloc::Join(_) => panic!("Item permanently borrowed/leaked"),
    }
  }
}

enum InternalOwned<T> {
  Item(T),
  Borrowed,
}

pub struct Owned<T>(InternalOwned<T>);
impl<T> Owned<T> {
  pub fn new(data:T) -> Self {
    Owned::<T>(InternalOwned::Item(data))
  }
  pub fn unwrap_or(self, other: T) -> T {
    if let InternalOwned::Item(x) = self.0 {
      x
    } else {
      other
    }
  }
  pub fn unwrap(self) -> T {
    if let InternalOwned::Item(x) = self.0 {
      x
    } else {
      panic!("Item permanently borrowed")
    }
  }
}

pub trait OwnedRetriever<SliceW: SliceWrapper<u8>+Send+'static> {
  fn view(&self) -> &[u8];
  fn unwrap(self) -> SliceW;
}

pub trait BatchSpawnable<T:Send+'static, Alloc:BrotliAlloc+Send+'static, SliceW:SliceWrapper<u8>+Send+'static>
  where <Alloc as Allocator<u8>>::AllocatedMemory:Send+'static
{
  type JoinHandle: Joinable<T, Alloc>;
  type FinalJoinHandle: OwnedRetriever<SliceW>;
  // this function takes in a SendAlloc per thread and converts them all into JoinHandle
  // the input is borrowed until the joins complete
  // owned is set to borrowed
  // the final join handle is a r/w lock which will return the SliceW to the owner
  // the FinalJoinHandle is only to be called when each individual JoinHandle has been examined
  fn batch_spawn<F: FnOnce(usize, &SliceW, Alloc) -> T>(
    &mut self,
    input: &mut Owned<SliceW>,
    alloc_per_thread:&mut [SendAlloc<T, Alloc, Self::JoinHandle>],
    f: F,
  ) -> Self::FinalJoinHandle;
}

pub fn CompressMultiSlice<Alloc:BrotliAlloc+Send+'static,
                          Spawner:BatchSpawnable<CompressionThreadResult<Alloc>,
                                                 Alloc,
                                                 <Alloc as Allocator<u8>>::AllocatedMemory>> (
  params:&BrotliEncoderParams,
  input_slice: &[u8],
  output: &mut [u8],
  alloc_per_thread:&mut [SendAlloc<CompressionThreadResult<Alloc>, Alloc, Spawner::JoinHandle>],
  mut thread_spawner: Spawner,
) -> Result<usize, ()> where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
  let input = if let InternalSendAlloc::A(ref mut alloc) = alloc_per_thread[0].0 {
    let mut input = <Alloc as Allocator<u8>>::alloc_cell(alloc, input_slice.len());
    input.slice_mut().clone_from_slice(input_slice);
    input
  } else {
    <Alloc as Allocator<u8>>::AllocatedMemory::default()
  };
  let mut owned_input = Owned::new(input);
  let ret = CompressMulti(params, &mut owned_input, output, alloc_per_thread,thread_spawner);
  if let InternalSendAlloc::A(ref mut alloc) = alloc_per_thread[0].0 {
    <Alloc as Allocator<u8>>::free_cell(alloc, owned_input.unwrap());
  }
  ret
}

fn get_range(thread_index: usize, num_threads: usize, file_size: usize) -> Range<usize> {
    ((thread_index * file_size) / num_threads)..(((thread_index + 1) * file_size) / num_threads)
}

pub fn CompressMulti<Alloc:BrotliAlloc+Send+'static,
                     SliceW: SliceWrapper<u8>+Send+'static,
                     Spawner:BatchSpawnable<CompressionThreadResult<Alloc>,
                                            Alloc,
                                            SliceW>> (
  params:&BrotliEncoderParams,
  owned_input: &mut Owned<SliceW>,
  output: &mut [u8],
  alloc_per_thread:&mut [SendAlloc<CompressionThreadResult<Alloc>, Alloc, Spawner::JoinHandle>],
  mut thread_spawner: Spawner,
) -> Result<usize, ()> where <Alloc as Allocator<u8>>::AllocatedMemory: Send {
    let num_threads = alloc_per_thread.len();
    let (mut alloc, alloc_rest) = alloc_per_thread.split_at_mut(1);
    let mut retrieve_owned_input = thread_spawner.batch_spawn(owned_input, alloc_rest, |_index,_input,mut alloc|{
        let mem = <Alloc as Allocator<u8>>::alloc_cell(&mut alloc, 0);
        CompressionThreadResult::<Alloc>{
            compressed:mem,
            compressed_size:0,
            alloc:alloc,
        }
    });
    let mut compression_result = Err(());
    let mut available_out = output.len();
    {
        let input = retrieve_owned_input.view();
        let mut state = BrotliEncoderCreateInstance(match mem::replace(&mut alloc[0].0,
                                                                             InternalSendAlloc::SpawningOrJoining(PhantomData::default())) {
            InternalSendAlloc::A(a) => a,
            _ => panic!("all public interfaces which create SendAlloc can only specify subtype A")
        });
        state.params = params.clone();
        state.params.appendable = true; // make sure we are at least appendable, so that future items can be catted in
        let mut range = get_range(0, num_threads, input.len());
        assert_eq!(range.start, 0);
        let mut out_offset = 0usize;
        while true {
            let mut next_in_offset = 0usize;
            let mut available_in = range.end - range.start;
            let result = BrotliEncoderCompressStream(&mut state,
                                                 BrotliEncoderOperation::BROTLI_OPERATION_FINISH,
                                                 &mut available_in,
                                                 &input[range.clone()],
                                                 &mut next_in_offset,
                                                 &mut available_out,
                                                 output,
                                                 &mut out_offset,
                                                 &mut None,
                                                 &mut |_a,_b,_c,_d|());
            let new_range = range.start + next_in_offset..range.end;
            range = new_range;
            if result != 0 {
                compression_result = Ok(out_offset);
                break;
            } else if available_out == 0 {
                compression_result = Err(()); // mark no space??
                break
            }
        }
        BrotliEncoderDestroyInstance(&mut state);
        alloc[0].0 = InternalSendAlloc::A(state.m8);
    }
    for thread in alloc_rest.iter_mut() {
        let alloc;
        match mem::replace(&mut thread.0, InternalSendAlloc::SpawningOrJoining(PhantomData::default())) {
            InternalSendAlloc::A(_) | InternalSendAlloc::SpawningOrJoining(_) => panic!("Thread not properly spawned"),
            InternalSendAlloc::Join(join) => match join.join() {
                Ok(result) => {
                    alloc = result.alloc;
                },
                Err(allocator) => {
                    alloc = allocator;
                    compression_result = Err(());
                },
            },
        }
        thread.0 = InternalSendAlloc::A(alloc);
    }
    *owned_input = Owned::new(retrieve_owned_input.unwrap()); // return the input to its rightful owner before returning
    compression_result
}

<|endoftext|>


<commit_before>extern crate lp_modeler;

use std::process;
use std::collections::HashMap;

use lp_modeler::variables::*;
use lp_modeler::variables::LpExpression::*;
use lp_modeler::operations::LpOperations;
use lp_modeler::problem::{LpObjective, LpProblem, LpFileFormat};
use lp_modeler::solvers::{SolverTrait, CbcSolver};

#[test]
fn test_readme_example_1() {
    let ref a = LpInteger::new("a");
    let ref b = LpInteger::new("b");
    let ref c = LpInteger::new("c");

    let mut problem = LpProblem::new("One Problem", LpObjective::Maximize);
    problem += 10.0 * a + 20.0 * b;

    problem += (500*a + 1200*b + 1500*c).le(10000);
    problem += (a).le(b);

    let solver = CbcSolver::new();

    match solver.run(&problem) {
        Ok((status, res)) => {
            println!("Status {:?}", status);
            for (name, value) in res.iter() {
                println!("value of {} = {}", name, value);
            }
        },
        Err(msg) => println!("{}", msg),
    }

    let output1 = "\\ One Problem

Maximize
  obj: 10 a + 20 b

Subject To
  c1: 500 a + 1200 b + 1500 c <= 10000
  c2: a - b <= 0

".to_string();
    let output2 = problem.to_lp_file_format();
    let output2 = output2.split("Generals").collect::<Vec<&str>>();
    let output2 = output2[0];
    assert_eq!(output1, output2);
}

#[test]
fn test_readme_example_2() {
    // Problem Data
    let men = vec!["A", "B", "C"];
    let women = vec!["D", "E", "F"];
    let scores: HashMap<(&'static str, &'static str), f32> = [
        (("A", "D"), 50.0),
        (("A", "E"), 75.0),
        (("A", "F"), 75.0),
        (("B", "D"), 60.0),
        (("B", "E"), 95.0),
        (("B", "F"), 80.0),
        (("C", "D"), 60.0),
        (("C", "E"), 70.0),
        (("C", "F"), 80.0),
        ].iter().cloned().collect();

    // Define Problem
    let mut problem = LpProblem::new("Matchmaking", LpObjective::Maximize);

    // Define Variables
    let mut vars = HashMap::new();
    for m in &men{
        for w in &women{
            vars.insert((m, w), LpBinary::new(&format!("{}_{}", m, w)));
        }
    }

    // Define Objective Function
    let mut obj_vec: Vec<LpExpression> = Vec::new();
    for (&(m, w), var) in &vars{
        let obj_coef = scores.get(&(m, w)).unwrap();
        obj_vec.push(*obj_coef * var);
    }
    problem += lp_sum(&obj_vec);

    // Define Constraints
    // Constraint 1: Each man must be assigned to exactly one woman
    for m in &men{
        let mut constr_vec = Vec::new();

        for w in &women{
            constr_vec.push(1.0 * vars.get(&(m, w)).unwrap());
        }

        problem += lp_sum(&constr_vec).equal(1);
    }

    // Constraint 2: Each woman must be assigned to exactly one man
    for w in &women{
        let mut constr_vec = Vec::new();

        for m in &men{
            constr_vec.push(1.0 * vars.get(&(m, w)).unwrap());
        }

        problem += lp_sum(&constr_vec).equal(1);
    }

    // Optionally write to file
    // let result = problem.write_lp("problem.lp");
    // match result{
    //     Ok(_) => println!("Written to file"),
    //     Err(msg) => println!("{}", msg)
    // }

    // Run Solver
    let solver = CbcSolver::new();
    let result = solver.run(&problem);

    // Terminate if error, or assign status & variable values
    assert!(result.is_ok(), result.unwrap_err());
    let (solver_status, var_values) = result.unwrap();

    // Compute final objective function value
    let mut obj_value = 0f32;
    for (&(m, w), var) in &vars{
        let obj_coef = scores.get(&(m, w)).unwrap();
        let var_value = var_values.get(&var.name).unwrap();

        obj_value += obj_coef * var_value;
    }

    // Print output
    println!("Status: {:?}", solver_status);
    println!("Objective Value: {}", obj_value);
    // println!("{:?}", var_values);
    for (var_name, var_value) in &var_values{
        let int_var_value = *var_value as u32;
        if int_var_value == 1{
            println!("{} = {}", var_name, int_var_value);
        }
    }

    assert_eq!(solver_status, lp_modeler::solvers::Status::Optimal);
    assert_eq!(obj_value, 230f32);
    assert_eq!(*var_values.get("A_F").unwrap(), 1f32);
    assert_eq!(*var_values.get("B_E").unwrap(), 1f32);
    assert_eq!(*var_values.get("C_D").unwrap(), 1f32);
}<commit_msg>Changed tuple signature for backwards compatibility<commit_after>extern crate lp_modeler;

use std::collections::HashMap;

use lp_modeler::variables::*;
use lp_modeler::variables::LpExpression::*;
use lp_modeler::operations::LpOperations;
use lp_modeler::problem::{LpObjective, LpProblem, LpFileFormat};
use lp_modeler::solvers::{SolverTrait, CbcSolver};

#[test]
fn test_readme_example_1() {
    let ref a = LpInteger::new("a");
    let ref b = LpInteger::new("b");
    let ref c = LpInteger::new("c");

    let mut problem = LpProblem::new("One Problem", LpObjective::Maximize);
    problem += 10.0 * a + 20.0 * b;

    problem += (500*a + 1200*b + 1500*c).le(10000);
    problem += (a).le(b);

    let solver = CbcSolver::new();

    match solver.run(&problem) {
        Ok((status, res)) => {
            println!("Status {:?}", status);
            for (name, value) in res.iter() {
                println!("value of {} = {}", name, value);
            }
        },
        Err(msg) => println!("{}", msg),
    }

    let output1 = "\\ One Problem

Maximize
  obj: 10 a + 20 b

Subject To
  c1: 500 a + 1200 b + 1500 c <= 10000
  c2: a - b <= 0

".to_string();
    let output2 = problem.to_lp_file_format();
    let output2 = output2.split("Generals").collect::<Vec<&str>>();
    let output2 = output2[0];
    assert_eq!(output1, output2);
}

#[test]
fn test_readme_example_2() {
    // Problem Data
    let men = vec!["A", "B", "C"];
    let women = vec!["D", "E", "F"];
    let scores: HashMap<(&'static str, &'static str), f32> = [
        (("A", "D"), 50.0),
        (("A", "E"), 75.0),
        (("A", "F"), 75.0),
        (("B", "D"), 60.0),
        (("B", "E"), 95.0),
        (("B", "F"), 80.0),
        (("C", "D"), 60.0),
        (("C", "E"), 70.0),
        (("C", "F"), 80.0),
        ].iter().cloned().collect();

    // Define Problem
    let mut problem = LpProblem::new("Matchmaking", LpObjective::Maximize);

    // Define Variables
    let mut vars = HashMap::new();
    for m in &men{
        for w in &women{
            vars.insert((m, w), LpBinary::new(&format!("{}_{}", m, w)));
        }
    }

    // Define Objective Function
    let mut obj_vec: Vec<LpExpression> = Vec::new();
    for (&(&m, &w), var) in &vars{
        let obj_coef = scores.get(&(m, w)).unwrap();
        obj_vec.push(*obj_coef * var);
    }
    problem += lp_sum(&obj_vec);

    // Define Constraints
    // Constraint 1: Each man must be assigned to exactly one woman
    for m in &men{
        let mut constr_vec = Vec::new();

        for w in &women{
            constr_vec.push(1.0 * vars.get(&(m, w)).unwrap());
        }

        problem += lp_sum(&constr_vec).equal(1);
    }

    // Constraint 2: Each woman must be assigned to exactly one man
    for w in &women{
        let mut constr_vec = Vec::new();

        for m in &men{
            constr_vec.push(1.0 * vars.get(&(m, w)).unwrap());
        }

        problem += lp_sum(&constr_vec).equal(1);
    }

    // Optionally write to file
    // let result = problem.write_lp("problem.lp");
    // match result{
    //     Ok(_) => println!("Written to file"),
    //     Err(msg) => println!("{}", msg)
    // }

    // Run Solver
    let solver = CbcSolver::new();
    let result = solver.run(&problem);

    // Terminate if error, or assign status & variable values
    assert!(result.is_ok(), result.unwrap_err());
    let (solver_status, var_values) = result.unwrap();

    // Compute final objective function value
    let mut obj_value = 0f32;
    for (&(&m, &w), var) in &vars{
        let obj_coef = scores.get(&(m, w)).unwrap();
        let var_value = var_values.get(&var.name).unwrap();

        obj_value += obj_coef * var_value;
    }

    // Print output
    println!("Status: {:?}", solver_status);
    println!("Objective Value: {}", obj_value);
    // println!("{:?}", var_values);
    for (var_name, var_value) in &var_values{
        let int_var_value = *var_value as u32;
        if int_var_value == 1{
            println!("{} = {}", var_name, int_var_value);
        }
    }

    assert_eq!(solver_status, lp_modeler::solvers::Status::Optimal);
    assert_eq!(obj_value, 230f32);
    assert_eq!(*var_values.get("A_F").unwrap(), 1f32);
    assert_eq!(*var_values.get("B_E").unwrap(), 1f32);
    assert_eq!(*var_values.get("C_D").unwrap(), 1f32);
}<|endoftext|>