<commit_before>#![feature(custom_derive)]

extern crate bincode;
extern crate ogg;
extern crate ogg_clock;
extern crate rustc_serialize;
extern crate serde;
extern crate ireul_interface;
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate byteorder;
extern crate url;
extern crate toml;
extern crate rand;

use std::thread;
use std::env;
use std::fmt;
use std::collections::HashSet;
use std::sync::mpsc::{self};
use std::net::{TcpStream, TcpListener};
use std::collections::{HashMap, VecDeque};
use std::io::{self, Read, Write};
use std::fs::File;

use byteorder::{ReadBytesExt, WriteBytesExt, BigEndian, ByteOrder};

use ogg::{OggTrack, OggTrackBuf, OggPage, OggPageBuf};
use ogg_clock::OggClock;
use ireul_interface::proxy::{
    SIZE_LIMIT,
    RequestWrapper,
    RequestType,
    BinderError,
    EnqueueTrackRequest,
    EnqueueTrackError,
    FastForward,
    FastForwardRequest,
    FastForwardError,
};

mod queue;
mod icecastwriter;

use queue::{PlayQueue, Handle, PlayQueueError};
use icecastwriter::{
    IceCastWriter,
    IceCastWriterOptions,
};

#[derive(RustcDecodable, Debug)]
struct MetadataConfig {
    name: Option<String>,
    description: Option<String>,
    url: Option<String>,
    genre: Option<String>,
}

#[derive(RustcDecodable, Debug)]
struct Config {
    icecast_url: String,
    metadata: Option<MetadataConfig>,
}

impl Config {
    fn icecast_writer_opts(&self) -> Result<IceCastWriterOptions, String> {
        let url = try!(url::Url::parse(&self.icecast_url)
            .map_err(|err| format!("Malformed URL: {:?}", err)));

        let mut opts = try!(IceCastWriterOptions::from_url(&url)
            .map_err(|err| format!("Unacceptable URL: {:?}", err)));

        if let Some(ref metadata) = self.metadata {
            if let Some(ref name) = metadata.name {
                opts.set_name(name);
            }
            if let Some(ref description) = metadata.description {
                opts.set_description(description);
            }
            if let Some(ref url) = metadata.url {
                opts.set_url(url);
            }
            if let Some(ref genre) = metadata.genre {
                opts.set_genre(genre);
            }
        }

        Ok(opts)
    }
}

fn main() {
    env_logger::init().unwrap();

    let config_file = env::args_os().nth(1).unwrap();
    let config: Config = {
        let mut reader = File::open(&config_file).expect("failed to open config file");
        let mut config_buf = String::new();
        reader.read_to_string(&mut config_buf).expect("failed to read config");
        toml::decode_str(&config_buf).expect("invalid config file")
    };
    let icecast_options = config.icecast_writer_opts().unwrap();

    let connector = IceCastWriter::new(icecast_options).unwrap();
    let mut file = File::open("howbigisthis.ogg").unwrap();
    let mut buffer = Vec::new();
    file.read_to_end(&mut buffer).unwrap();
    let offline_track = OggTrackBuf::new(buffer).unwrap();

    let play_queue = PlayQueue::new(50);

    let output_manager = OutputManager {
        connector: connector,
        cur_serial: 0,
        cur_sequence: 0,
        // position: 0,
        clock: OggClock::new(48000),
        playing_offline: false,
        buffer: VecDeque::new(),
        play_queue: PlayQueue::new(10),
        offline_track: offline_track,
    };

    let control = TcpListener::bind("0.0.0.0:3001").unwrap();
    let mut core = Core::new(control, output_manager).unwrap();
    loop {
        core.tick();

    }
}

fn validate_positions(track: &OggTrack) -> Result<(), ()> {
    let mut current = 0;
    let mut is_first = true;

    for page in track.pages() {
        let position = page.position();

        if is_first {
            is_first = false;

            if position != 0 {
                return Err(());
            }
        }

        if position < current {
            return Err(());
        }
        current = position;
    }

    Ok(())
}

fn check_sample_rate(req: u32, track: &OggTrack) -> Result<(), ()> {
    warn!("check_sample_rate: STUB");
    Ok(())
}

fn update_serial(serial: u32, track: &mut OggTrack) {
    for page in track.pages_mut() {
        page.set_serial(serial);
    }
}

fn update_positions(start_pos: u64, track: &mut OggTrack) {
    for page in track.pages_mut() {
        let old_pos = page.position();
        page.set_position(start_pos + old_pos);
    }
}

fn final_position(track: &OggTrack) -> Option<u64> {
    let mut position = None;
    for page in track.pages() {
        position = Some(page.position());
    }
    position
}

struct Core {
    output: OutputManager,
    proxy_rx: mpsc::Receiver<RequestWrapper>,
}

impl Core {
    fn new(control: TcpListener, om: OutputManager) -> io::Result<Core> {
        let (tx, rx) = mpsc::sync_channel(5);

        let proxy_tx_client = tx.clone();
        thread::spawn(move || {
            client_acceptor(control, proxy_tx_client);
        });

        Ok(Core {
            output: om,
            proxy_rx: rx,
        })
    }

    fn enqueue_track(&mut self, req: EnqueueTrackRequest) -> Result<Handle, EnqueueTrackError> {
        let EnqueueTrackRequest { mut track } = req;
        {
            let mut pages = 0;
            let mut samples = 0;
            for page in track.pages() {
                pages += 1;
                samples = page.position();
            }

            info!("a client sent {} samples in {} pages", samples, pages);
        }
        if track.as_u8_slice().len() == 0 {
            return Err(EnqueueTrackError::InvalidTrack);
        }

        try!(validate_positions(&track)
            .map_err(|()| EnqueueTrackError::InvalidTrack));

        try!(check_sample_rate(self.output.clock.sample_rate(), &track)
            .map_err(|()| EnqueueTrackError::BadSampleRate));

        self.output.play_queue.add_track(track.as_ref())
            .map_err(|err| match err {
                PlayQueueError::Full => EnqueueTrackError::Full,
            })
    }

    fn fast_forward(&mut self, req: FastForwardRequest) -> Result<(), FastForwardError> {
        try!(self.output.fast_forward(req.kind));
        Ok(())
    }

    fn handle_command(&mut self, req_wr: RequestWrapper) {
        let mut binder = CoreBinder { core: self };
        binder.handle_command(req_wr)
    }

    fn tick(&mut self) {
        loop {
            match self.proxy_rx.try_recv() {
                Ok(cmd) => self.handle_command(cmd),
                Err(err) => break,
            }
        }
        self.output.copy_page();
    }
}

fn client_worker(mut stream: TcpStream, chan: mpsc::SyncSender<RequestWrapper>) -> io::Result<()> {
    const BUFFER_SIZE_LIMIT: usize = 20 * 1 << 20;
    loop {
        let version = try!(stream.read_u8());

        if version != 0 {
            let err_msg = format!("invalid version: {}", version);
            return Err(io::Error::new(io::ErrorKind::Other, err_msg));
        }

        let op_code = try!(stream.read_u32::<BigEndian>());
        if op_code == 0 {
            info!("goodbye, client");
            return Ok(());
        }

        let req_type = try!(RequestType::from_op_code(op_code).map_err(|_| {
            let err_msg = format!("unknown op-code {:?}", op_code);
            io::Error::new(io::ErrorKind::Other, err_msg)
        }));

        let frame_length = try!(stream.read_u32::<BigEndian>()) as usize;
        if BUFFER_SIZE_LIMIT < frame_length {
            let err_msg = format!("datagram too large: {} bytes (limit is {})",
                frame_length, BUFFER_SIZE_LIMIT);
            return Err(io::Error::new(io::ErrorKind::Other, err_msg));
        }

        let mut req_buf = Vec::new();
        {
            let mut limit_reader = Read::by_ref(&mut stream).take(frame_length as u64);
            try!(limit_reader.read_to_end(&mut req_buf));
        }

        if req_buf.len() != frame_length {
            let err_msg = format!(
                "datagram truncated: got {} bytes, expected {}",
                req_buf.len(), frame_length);
            return Err(io::Error::new(io::ErrorKind::Other, err_msg));
        }

        let (resp_tx, resp_rx) = mpsc::sync_channel(1);
        chan.send(RequestWrapper {
            response_queue: resp_tx,
            req_type: req_type,
            req_buf: req_buf,
        }).unwrap();

        let response = resp_rx.recv().unwrap();
        try!(stream.write_u32::<BigEndian>(response.len() as u32));
        try!(stream.write_all(&response));
    }
}

fn client_acceptor(server: TcpListener, chan: mpsc::SyncSender<RequestWrapper>) {
    for stream in server.incoming() {
        match stream {
            Ok(stream) => {
                let client_chan = chan.clone();
                thread::spawn(move || {
                    if let Err(err) = client_worker(stream, client_chan) {
                        info!("client disconnected with error: {:?}", err);
                    }
                });
            },
            Err(err) => {
                info!("error accepting new client: {:?}", err);
            }
        }
    }
}

struct CoreBinder<'a> {
    core: &'a mut Core,
}

impl<'a> CoreBinder<'a> {
    fn handle_command(&mut self, req_wr: RequestWrapper) {
        info!("CoreBinder::handle_command");
        let RequestWrapper {
            response_queue: response_queue,
            req_type: req_type,
            req_buf: req_buf,
        } = req_wr;
        let response = match req_type {
            RequestType::EnqueueTrack => {
                self.enqueue_track(&req_buf)
            },
            RequestType::FastForward => {
                self.fast_forward(&req_buf)
            },
        };
        response_queue.send(response).unwrap();
    }

    fn enqueue_track(&mut self, req: &[u8]) -> Vec<u8> {
        info!("CoreBinder::enqueue_track");
        let track_res: Result<Handle, EnqueueTrackError> =
            OggTrack::new(req)
                .map(|t| EnqueueTrackRequest { track: t.to_owned() })
                .map_err(|_| EnqueueTrackError::InvalidTrack)
                .and_then(|req| self.core.enqueue_track(req));

        let mut buf = io::Cursor::new(Vec::new());
        match track_res {
            Ok(Handle(handle_id)) => {
                buf.write_u8(0).unwrap();
                buf.write_u64::<BigEndian>(handle_id).unwrap();
            },
            Err(err) => {
                buf.write_u8(1).unwrap();
                // error kind
                buf.write_u32::<BigEndian>(err as u32).unwrap();
                // future expansion: error message
                buf.write_u32::<BigEndian>(0).unwrap();
            }
        }
        buf.into_inner()
    }

    fn fast_forward(&mut self, req: &[u8]) -> Vec<u8> {
        info!("CoreBinder::fast_forward");
        if req.len() < 4 {
            // should be an error condition... ProtocolError.
            // we need to disconnect the client in this case, usually.
            return Vec::new();
        }

        let ff_type = BigEndian::read_u32(req);
        let req: FastForwardRequest =
            FastForward::from_u32(ff_type)
                .map(|ff| FastForwardRequest { kind: ff })
                .unwrap();

        let mut buf = io::Cursor::new(Vec::new());
        info!("core.fast_forward({:?})", req);
        match self.core.fast_forward(req) {
            Ok(()) => {
                buf.write_u8(0).unwrap();
            },
            Err(err) => {
                buf.write_u8(1).unwrap();
                // error kind
                buf.write_u32::<BigEndian>(0).unwrap();
                // future expansion: error message
                buf.write_u32::<BigEndian>(0).unwrap();
            }
        };
        buf.into_inner()
    }
}

/// Connects to IceCast and holds references to streamable content.
struct OutputManager {
    connector: IceCastWriter,
    cur_serial: u32,
    cur_sequence: u32,
    clock: OggClock,

    playing_offline: bool,
    buffer: VecDeque<OggPageBuf>,
    play_queue: PlayQueue,
    offline_track: OggTrackBuf,
}

impl OutputManager {
    fn fill_buffer(&mut self) {
        let mut track = match self.play_queue.pop_track() {
            Some(track) => {
                self.playing_offline = false;
                track
            },
            None => {
                self.playing_offline = true;
                self.offline_track.clone()
            }
        };

        // not sure why we as_mut instead of just using &mut track
        update_serial(self.cur_serial, track.as_mut());
        self.cur_serial = self.cur_serial.wrapping_add(0);

        self.buffer.extend(track.pages().map(|x| x.to_owned()));
    }

    fn get_next_page(&mut self) -> OggPageBuf {
        if self.buffer.is_empty() {
            self.fill_buffer();
        }
        self.buffer.pop_front().unwrap()
    }

    fn fast_forward_track_boundary(&mut self) -> Result<(), FastForwardError> {
        loop {
            let page = self.get_next_page();
            debug!("checking page...");
            if page_starts_track(page.as_ref()) {
                debug!("checking page... found a start");
                self.buffer.push_front(page);
                break;
            }
        }
        Ok(())
    }

    fn fast_forward(&mut self, kind: FastForward) -> Result<(), FastForwardError> {
        match kind {
            FastForward::TrackBoundary => {
                self.fast_forward_track_boundary()
            }
        }
    }

    fn copy_page(&mut self) {
        let page = self.get_next_page();
        self.clock.wait(&page).unwrap();
        self.connector.send_ogg_page(&page).unwrap();

        debug!("copied page :: granule_pos = {:?}; serial = {:?}; sequence = {:?}",
            page.position(),
            page.serial(),
            page.sequence());
    }
}

fn page_starts_track(page: &OggPage) -> bool {
    page.body().starts_with(b"\x01vorbis")
}
<commit_msg>Automatically skip offline track when a real song is enqueued<commit_after>#![feature(custom_derive)]

extern crate bincode;
extern crate ogg;
extern crate ogg_clock;
extern crate rustc_serialize;
extern crate serde;
extern crate ireul_interface;
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate byteorder;
extern crate url;
extern crate toml;
extern crate rand;

use std::thread;
use std::env;
use std::fmt;
use std::collections::HashSet;
use std::sync::mpsc::{self};
use std::net::{TcpStream, TcpListener};
use std::collections::{HashMap, VecDeque};
use std::io::{self, Read, Write};
use std::fs::File;

use byteorder::{ReadBytesExt, WriteBytesExt, BigEndian, ByteOrder};

use ogg::{OggTrack, OggTrackBuf, OggPage, OggPageBuf};
use ogg_clock::OggClock;
use ireul_interface::proxy::{
    SIZE_LIMIT,
    RequestWrapper,
    RequestType,
    BinderError,
    EnqueueTrackRequest,
    EnqueueTrackError,
    FastForward,
    FastForwardRequest,
    FastForwardError,
};

mod queue;
mod icecastwriter;

use queue::{PlayQueue, Handle, PlayQueueError};
use icecastwriter::{
    IceCastWriter,
    IceCastWriterOptions,
};

#[derive(RustcDecodable, Debug)]
struct MetadataConfig {
    name: Option<String>,
    description: Option<String>,
    url: Option<String>,
    genre: Option<String>,
}

#[derive(RustcDecodable, Debug)]
struct Config {
    icecast_url: String,
    metadata: Option<MetadataConfig>,
}

impl Config {
    fn icecast_writer_opts(&self) -> Result<IceCastWriterOptions, String> {
        let url = try!(url::Url::parse(&self.icecast_url)
            .map_err(|err| format!("Malformed URL: {:?}", err)));

        let mut opts = try!(IceCastWriterOptions::from_url(&url)
            .map_err(|err| format!("Unacceptable URL: {:?}", err)));

        if let Some(ref metadata) = self.metadata {
            if let Some(ref name) = metadata.name {
                opts.set_name(name);
            }
            if let Some(ref description) = metadata.description {
                opts.set_description(description);
            }
            if let Some(ref url) = metadata.url {
                opts.set_url(url);
            }
            if let Some(ref genre) = metadata.genre {
                opts.set_genre(genre);
            }
        }

        Ok(opts)
    }
}

fn main() {
    env_logger::init().unwrap();

    let config_file = env::args_os().nth(1).unwrap();
    let config: Config = {
        let mut reader = File::open(&config_file).expect("failed to open config file");
        let mut config_buf = String::new();
        reader.read_to_string(&mut config_buf).expect("failed to read config");
        toml::decode_str(&config_buf).expect("invalid config file")
    };
    let icecast_options = config.icecast_writer_opts().unwrap();

    let connector = IceCastWriter::new(icecast_options).unwrap();
    let mut file = File::open("howbigisthis.ogg").unwrap();
    let mut buffer = Vec::new();
    file.read_to_end(&mut buffer).unwrap();
    let offline_track = OggTrackBuf::new(buffer).unwrap();

    let play_queue = PlayQueue::new(50);

    let output_manager = OutputManager {
        connector: connector,
        cur_serial: 0,
        cur_sequence: 0,
        // position: 0,
        clock: OggClock::new(48000),
        playing_offline: false,
        buffer: VecDeque::new(),
        play_queue: PlayQueue::new(10),
        offline_track: offline_track,
    };

    let control = TcpListener::bind("0.0.0.0:3001").unwrap();
    let mut core = Core::new(control, output_manager).unwrap();
    loop {
        core.tick();

    }
}

fn validate_positions(track: &OggTrack) -> Result<(), ()> {
    let mut current = 0;
    let mut is_first = true;

    for page in track.pages() {
        let position = page.position();

        if is_first {
            is_first = false;

            if position != 0 {
                return Err(());
            }
        }

        if position < current {
            return Err(());
        }
        current = position;
    }

    Ok(())
}

fn check_sample_rate(req: u32, track: &OggTrack) -> Result<(), ()> {
    warn!("check_sample_rate: STUB");
    Ok(())
}

fn update_serial(serial: u32, track: &mut OggTrack) {
    for page in track.pages_mut() {
        page.set_serial(serial);
    }
}

fn update_positions(start_pos: u64, track: &mut OggTrack) {
    for page in track.pages_mut() {
        let old_pos = page.position();
        page.set_position(start_pos + old_pos);
    }
}

fn final_position(track: &OggTrack) -> Option<u64> {
    let mut position = None;
    for page in track.pages() {
        position = Some(page.position());
    }
    position
}

struct Core {
    output: OutputManager,
    proxy_rx: mpsc::Receiver<RequestWrapper>,
}

impl Core {
    fn new(control: TcpListener, om: OutputManager) -> io::Result<Core> {
        let (tx, rx) = mpsc::sync_channel(5);

        let proxy_tx_client = tx.clone();
        thread::spawn(move || {
            client_acceptor(control, proxy_tx_client);
        });

        Ok(Core {
            output: om,
            proxy_rx: rx,
        })
    }

    fn enqueue_track(&mut self, req: EnqueueTrackRequest) -> Result<Handle, EnqueueTrackError> {
        let EnqueueTrackRequest { mut track } = req;
        {
            let mut pages = 0;
            let mut samples = 0;
            for page in track.pages() {
                pages += 1;
                samples = page.position();
            }

            info!("a client sent {} samples in {} pages", samples, pages);
        }
        if track.as_u8_slice().len() == 0 {
            return Err(EnqueueTrackError::InvalidTrack);
        }

        try!(validate_positions(&track)
            .map_err(|()| EnqueueTrackError::InvalidTrack));

        try!(check_sample_rate(self.output.clock.sample_rate(), &track)
            .map_err(|()| EnqueueTrackError::BadSampleRate));

        let handle = self.output.play_queue.add_track(track.as_ref())
            .map_err(|err| match err {
                PlayQueueError::Full => EnqueueTrackError::Full,
            });

        if self.output.playing_offline {
            self.output.fast_forward_track_boundary();
        }

        handle
    }

    fn fast_forward(&mut self, req: FastForwardRequest) -> Result<(), FastForwardError> {
        try!(self.output.fast_forward(req.kind));
        Ok(())
    }

    fn handle_command(&mut self, req_wr: RequestWrapper) {
        let mut binder = CoreBinder { core: self };
        binder.handle_command(req_wr)
    }

    fn tick(&mut self) {
        loop {
            match self.proxy_rx.try_recv() {
                Ok(cmd) => self.handle_command(cmd),
                Err(err) => break,
            }
        }
        self.output.copy_page();
    }
}

fn client_worker(mut stream: TcpStream, chan: mpsc::SyncSender<RequestWrapper>) -> io::Result<()> {
    const BUFFER_SIZE_LIMIT: usize = 20 * 1 << 20;
    loop {
        let version = try!(stream.read_u8());

        if version != 0 {
            let err_msg = format!("invalid version: {}", version);
            return Err(io::Error::new(io::ErrorKind::Other, err_msg));
        }

        let op_code = try!(stream.read_u32::<BigEndian>());
        if op_code == 0 {
            info!("goodbye, client");
            return Ok(());
        }

        let req_type = try!(RequestType::from_op_code(op_code).map_err(|_| {
            let err_msg = format!("unknown op-code {:?}", op_code);
            io::Error::new(io::ErrorKind::Other, err_msg)
        }));

        let frame_length = try!(stream.read_u32::<BigEndian>()) as usize;
        if BUFFER_SIZE_LIMIT < frame_length {
            let err_msg = format!("datagram too large: {} bytes (limit is {})",
                frame_length, BUFFER_SIZE_LIMIT);
            return Err(io::Error::new(io::ErrorKind::Other, err_msg));
        }

        let mut req_buf = Vec::new();
        {
            let mut limit_reader = Read::by_ref(&mut stream).take(frame_length as u64);
            try!(limit_reader.read_to_end(&mut req_buf));
        }

        if req_buf.len() != frame_length {
            let err_msg = format!(
                "datagram truncated: got {} bytes, expected {}",
                req_buf.len(), frame_length);
            return Err(io::Error::new(io::ErrorKind::Other, err_msg));
        }

        let (resp_tx, resp_rx) = mpsc::sync_channel(1);
        chan.send(RequestWrapper {
            response_queue: resp_tx,
            req_type: req_type,
            req_buf: req_buf,
        }).unwrap();

        let response = resp_rx.recv().unwrap();
        try!(stream.write_u32::<BigEndian>(response.len() as u32));
        try!(stream.write_all(&response));
    }
}

fn client_acceptor(server: TcpListener, chan: mpsc::SyncSender<RequestWrapper>) {
    for stream in server.incoming() {
        match stream {
            Ok(stream) => {
                let client_chan = chan.clone();
                thread::spawn(move || {
                    if let Err(err) = client_worker(stream, client_chan) {
                        info!("client disconnected with error: {:?}", err);
                    }
                });
            },
            Err(err) => {
                info!("error accepting new client: {:?}", err);
            }
        }
    }
}

struct CoreBinder<'a> {
    core: &'a mut Core,
}

impl<'a> CoreBinder<'a> {
    fn handle_command(&mut self, req_wr: RequestWrapper) {
        info!("CoreBinder::handle_command");
        let RequestWrapper {
            response_queue: response_queue,
            req_type: req_type,
            req_buf: req_buf,
        } = req_wr;
        let response = match req_type {
            RequestType::EnqueueTrack => {
                self.enqueue_track(&req_buf)
            },
            RequestType::FastForward => {
                self.fast_forward(&req_buf)
            },
        };
        response_queue.send(response).unwrap();
    }

    fn enqueue_track(&mut self, req: &[u8]) -> Vec<u8> {
        info!("CoreBinder::enqueue_track");
        let track_res: Result<Handle, EnqueueTrackError> =
            OggTrack::new(req)
                .map(|t| EnqueueTrackRequest { track: t.to_owned() })
                .map_err(|_| EnqueueTrackError::InvalidTrack)
                .and_then(|req| self.core.enqueue_track(req));

        let mut buf = io::Cursor::new(Vec::new());
        match track_res {
            Ok(Handle(handle_id)) => {
                buf.write_u8(0).unwrap();
                buf.write_u64::<BigEndian>(handle_id).unwrap();
            },
            Err(err) => {
                buf.write_u8(1).unwrap();
                // error kind
                buf.write_u32::<BigEndian>(err as u32).unwrap();
                // future expansion: error message
                buf.write_u32::<BigEndian>(0).unwrap();
            }
        }
        buf.into_inner()
    }

    fn fast_forward(&mut self, req: &[u8]) -> Vec<u8> {
        info!("CoreBinder::fast_forward");
        if req.len() < 4 {
            // should be an error condition... ProtocolError.
            // we need to disconnect the client in this case, usually.
            return Vec::new();
        }

        let ff_type = BigEndian::read_u32(req);
        let req: FastForwardRequest =
            FastForward::from_u32(ff_type)
                .map(|ff| FastForwardRequest { kind: ff })
                .unwrap();

        let mut buf = io::Cursor::new(Vec::new());
        info!("core.fast_forward({:?})", req);
        match self.core.fast_forward(req) {
            Ok(()) => {
                buf.write_u8(0).unwrap();
            },
            Err(err) => {
                buf.write_u8(1).unwrap();
                // error kind
                buf.write_u32::<BigEndian>(0).unwrap();
                // future expansion: error message
                buf.write_u32::<BigEndian>(0).unwrap();
            }
        };
        buf.into_inner()
    }
}

/// Connects to IceCast and holds references to streamable content.
struct OutputManager {
    connector: IceCastWriter,
    cur_serial: u32,
    cur_sequence: u32,
    clock: OggClock,

    playing_offline: bool,
    buffer: VecDeque<OggPageBuf>,
    play_queue: PlayQueue,
    offline_track: OggTrackBuf,
}

impl OutputManager {
    fn fill_buffer(&mut self) {
        let mut track = match self.play_queue.pop_track() {
            Some(track) => {
                self.playing_offline = false;
                track
            },
            None => {
                self.playing_offline = true;
                self.offline_track.clone()
            }
        };

        // not sure why we as_mut instead of just using &mut track
        update_serial(self.cur_serial, track.as_mut());
        self.cur_serial = self.cur_serial.wrapping_add(0);

        self.buffer.extend(track.pages().map(|x| x.to_owned()));
    }

    fn get_next_page(&mut self) -> OggPageBuf {
        if self.buffer.is_empty() {
            self.fill_buffer();
        }
        self.buffer.pop_front().unwrap()
    }

    fn fast_forward_track_boundary(&mut self) -> Result<(), FastForwardError> {
        loop {
            let page = self.get_next_page();
            debug!("checking page...");
            if page_starts_track(page.as_ref()) {
                debug!("checking page... found a start");
                self.buffer.push_front(page);
                break;
            }
        }
        Ok(())
    }

    fn fast_forward(&mut self, kind: FastForward) -> Result<(), FastForwardError> {
        match kind {
            FastForward::TrackBoundary => {
                self.fast_forward_track_boundary()
            }
        }
    }

    fn copy_page(&mut self) {
        let page = self.get_next_page();
        self.clock.wait(&page).unwrap();
        self.connector.send_ogg_page(&page).unwrap();

        debug!("copied page :: granule_pos = {:?}; serial = {:?}; sequence = {:?}",
            page.position(),
            page.serial(),
            page.sequence());
    }
}

fn page_starts_track(page: &OggPage) -> bool {
    page.body().starts_with(b"\x01vorbis")
}
<|endoftext|>


<commit_before>pub mod api;
pub mod http;

use serde_json;
use std;
use std::fmt;

#[derive(Debug)]
pub enum Error {
	Api(api::Api),
	Http,
	Json(serde_json::Error),
	Unknown,
}

impl Error {
	pub fn is_api_error(self) -> bool {
		match self {
			Error::Api(_) => true,
			Error::Json(_) => false,
			Error::Http => false,
			_ => false,
		}
	}

	pub fn is_http_err(self) -> bool {
		match self {
			Error::Api(_) => false,
			Error::Json(_) => false,
			Error::Http => true,
			_ => false,
		}
	}

	pub fn is_json_error(self) -> bool {
		match self {
			Error::Api(_) => false,
			Error::Json(_) => true,
			Error::Http => false,
			_ => false,
		}
	}
}

impl fmt::Display for Error {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		match *self {
			Error::Api(ref e) => e.fmt(f),
			Error::Http => write!(f, "HTTP"),
			Error::Json(_) => write!(f, "Json"),
			Error::Unknown => write!(f, "Unknown"),
		}
	}
}

impl std::error::Error for Error {
	fn description(&self) -> &str {
		match *self {
			Error::Api(_) => "API error",
			Error::Http => "a connection error",
			Error::Json(_) => "Json error",
			Error::Unknown => "an unknown error",
		}
	}
}
<commit_msg>pass json errors to the application<commit_after>pub mod api;
pub mod http;

use serde_json;
use std;
use std::fmt;

#[derive(Debug)]
pub enum Error {
	Api(api::Api),
	Http,
	Json(serde_json::Error),
	Unknown,
}

impl Error {
	pub fn is_api_error(self) -> bool {
		match self {
			Error::Api(_) => true,
			Error::Json(_) => false,
			Error::Http => false,
			_ => false,
		}
	}

	pub fn is_http_err(self) -> bool {
		match self {
			Error::Api(_) => false,
			Error::Json(_) => false,
			Error::Http => true,
			_ => false,
		}
	}

	pub fn is_json_error(self) -> bool {
		match self {
			Error::Api(_) => false,
			Error::Json(_) => true,
			Error::Http => false,
			_ => false,
		}
	}
}

impl fmt::Display for Error {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		match *self {
			Error::Api(ref e) => e.fmt(f),
			Error::Http => write!(f, "HTTP"),
			Error::Json(ref e) => write!(f, "Json: {}", e),
			Error::Unknown => write!(f, "Unknown"),
		}
	}
}

impl std::error::Error for Error {
	fn description(&self) -> &str {
		match *self {
			Error::Api(_) => "API error",
			Error::Http => "a connection error",
			Error::Json(_) => "Json error",
			Error::Unknown => "an unknown error",
		}
	}
}
<|endoftext|>


<commit_before>use super::{Function, Parameterised, VFunction, QFunction, Projector, Projection};
use geometry::Space;
use ndarray::{ArrayView, Array2};
use std::marker::PhantomData;

#[derive(Serialize, Deserialize)]
pub struct Linear<S: Space, P: Projector<S>> {
    pub projector: P,
    pub weights: Array2<f64>,

    phantom: PhantomData<S>,
}

impl<S: Space, P: Projector<S>> Linear<S, P> {
    pub fn new(projector: P, n_outputs: usize) -> Self {
        let n_features = projector.size();

        Self {
            projector: projector,
            weights: Array2::<f64>::zeros((n_features, n_outputs)),

            phantom: PhantomData,
        }
    }
}

impl<S: Space, P: Projector<S>> Function<S::Repr, f64> for Linear<S, P> {
    fn evaluate(&self, input: &S::Repr) -> f64 {
        // Compute the feature vector phi:
        let phi = self.projector.project(input);

        <Self as VFunction<S>>::evaluate_phi(self, &phi)
    }
}

impl<S: Space, P: Projector<S>> Function<S::Repr, Vec<f64>> for Linear<S, P> {
    fn evaluate(&self, input: &S::Repr) -> Vec<f64> {
        // Compute the feature vector phi:
        let phi = self.projector.project(input);

        // Apply matrix multiplication and return Vec<f64>:
        <Self as QFunction<S>>::evaluate_phi(self, &phi)
    }
}

impl<S: Space, P: Projector<S>> Parameterised<S::Repr, f64> for Linear<S, P> {
    fn update(&mut self, input: &S::Repr, error: f64) {
        // Compute the feature vector phi:
        let phi = self.projector.project(input);

        <Self as VFunction<S>>::update_phi(self, &phi, error);
    }
}

impl<S: Space, P: Projector<S>> Parameterised<S::Repr, Vec<f64>> for Linear<S, P> {
    fn update(&mut self, input: &S::Repr, errors: Vec<f64>) {
        // Compute the feature vector phi:
        let phi = self.projector.project(input);

        <Self as QFunction<S>>::update_phi(self, &phi, errors);
    }
}

impl<S: Space, P: Projector<S>> VFunction<S> for Linear<S, P> {
    fn evaluate_phi(&self, phi: &Projection) -> f64 {
        <Self as QFunction<S>>::evaluate_action_phi(self, phi, 0)
    }

    fn update_phi(&mut self, phi: &Projection, error: f64) {
        <Self as QFunction<S>>::update_action_phi(self, phi, 0, error);
    }
}

impl<S: Space, P: Projector<S>> QFunction<S> for Linear<S, P> {
    fn evaluate_action(&self, input: &S::Repr, action: usize) -> f64 {
        let phi = self.projector.project(input);

        self.evaluate_action_phi(&phi, action)
    }

    fn update_action(&mut self, input: &S::Repr, action: usize, error: f64) {
        let phi = self.projector.project(input);

        self.update_action_phi(&phi, action, error);
    }

    fn evaluate_phi(&self, phi: &Projection) -> Vec<f64> {
        match phi {
            &Projection::Dense(ref dense_phi) => (self.weights.t().dot(dense_phi)).into_raw_vec(),
            &Projection::Sparse(ref sparse_phi) =>
                (0..self.weights.cols()).map(|c| {
                    sparse_phi.iter().fold(0.0, |acc, idx| acc + self.weights[(*idx, c)])
                }).collect(),
        }
    }

    fn evaluate_action_phi(&self, phi: &Projection, action: usize) -> f64 {
        let col = self.weights.column(action);

        match phi {
            &Projection::Dense(ref dense_phi) => col.dot(&(dense_phi/phi.z())),
            &Projection::Sparse(ref sparse_phi) =>
                sparse_phi.iter().fold(0.0, |acc, idx| acc + col[*idx]),
        }
    }

    fn update_phi(&mut self, phi: &Projection, errors: Vec<f64>) {
        let z = phi.z();
        let sf = 1.0/z;

        match phi {
            &Projection::Dense(ref dense_phi) => {
                let view = dense_phi.view().into_shape((self.weights.rows(), 1)).unwrap();
                let error_matrix = ArrayView::from_shape((1, self.weights.cols()), errors.as_slice())
                    .unwrap();

                self.weights.scaled_add(sf, &view.dot(&error_matrix))
            },
            &Projection::Sparse(ref sparse_phi) => {
                for c in 0..self.weights.cols() {
                    let mut col = self.weights.column_mut(c);
                    let error = errors[c];

                    for idx in sparse_phi {
                        col[*idx] += error
                    }
                }
            },
        }
    }

    fn update_action_phi(&mut self, phi: &Projection, action: usize, error: f64) {
        let mut col = self.weights.column_mut(action);
        let scaled_error = error/phi.z();

        match phi {
            &Projection::Dense(ref dense_phi) => col.scaled_add(scaled_error, dense_phi),
            &Projection::Sparse(ref sparse_phi) => {
                for idx in sparse_phi {
                    col[*idx] += scaled_error
                }
            },
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate seahash;

    use super::*;
    use std::hash::BuildHasherDefault;
    use fa::projection::TileCoding;

    type SHBuilder = BuildHasherDefault<seahash::SeaHasher>;

    #[test]
    fn test_dense_update_eval() {
        let p = TileCoding::new(SHBuilder::default(), 4, 100);
        let mut f = Linear::new(p.clone(), 1);

        let input = vec![5.0];

        f.update(&input, 50.0);
        let out: f64 = f.evaluate(&input);

        assert!(out > 0.0);
    }

    #[test]
    fn test_dense_eval_phi() {
        let p = TileCoding::new(SHBuilder::default(), 4, 100);
        let f = Linear::new(p.clone(), 1);

        let input = vec![5.0];

        let out: f64 = f.evaluate(&input);
        let out_alt_v: f64 = VFunction::evaluate_phi(&f, &p.project(&input));
        let out_alt_q: Vec<f64> = QFunction::evaluate_phi(&f, &p.project(&input));

        assert_eq!(out, out_alt_v);
        assert_eq!(vec![out], out_alt_q);
    }
}
<commit_msg>fix(fa): fix `Linear::evaluate_action_phi`<commit_after>use super::{Function, Parameterised, VFunction, QFunction, Projector, Projection};
use geometry::Space;
use ndarray::{ArrayView, Array2};
use std::marker::PhantomData;

#[derive(Serialize, Deserialize)]
pub struct Linear<S: Space, P: Projector<S>> {
    pub projector: P,
    pub weights: Array2<f64>,

    phantom: PhantomData<S>,
}

impl<S: Space, P: Projector<S>> Linear<S, P> {
    pub fn new(projector: P, n_outputs: usize) -> Self {
        let n_features = projector.size();

        Self {
            projector: projector,
            weights: Array2::<f64>::zeros((n_features, n_outputs)),

            phantom: PhantomData,
        }
    }
}

impl<S: Space, P: Projector<S>> Function<S::Repr, f64> for Linear<S, P> {
    fn evaluate(&self, input: &S::Repr) -> f64 {
        // Compute the feature vector phi:
        let phi = self.projector.project(input);

        <Self as VFunction<S>>::evaluate_phi(self, &phi)
    }
}

impl<S: Space, P: Projector<S>> Function<S::Repr, Vec<f64>> for Linear<S, P> {
    fn evaluate(&self, input: &S::Repr) -> Vec<f64> {
        // Compute the feature vector phi:
        let phi = self.projector.project(input);

        // Apply matrix multiplication and return Vec<f64>:
        <Self as QFunction<S>>::evaluate_phi(self, &phi)
    }
}

impl<S: Space, P: Projector<S>> Parameterised<S::Repr, f64> for Linear<S, P> {
    fn update(&mut self, input: &S::Repr, error: f64) {
        // Compute the feature vector phi:
        let phi = self.projector.project(input);

        <Self as VFunction<S>>::update_phi(self, &phi, error);
    }
}

impl<S: Space, P: Projector<S>> Parameterised<S::Repr, Vec<f64>> for Linear<S, P> {
    fn update(&mut self, input: &S::Repr, errors: Vec<f64>) {
        // Compute the feature vector phi:
        let phi = self.projector.project(input);

        <Self as QFunction<S>>::update_phi(self, &phi, errors);
    }
}

impl<S: Space, P: Projector<S>> VFunction<S> for Linear<S, P> {
    fn evaluate_phi(&self, phi: &Projection) -> f64 {
        <Self as QFunction<S>>::evaluate_action_phi(self, phi, 0)
    }

    fn update_phi(&mut self, phi: &Projection, error: f64) {
        <Self as QFunction<S>>::update_action_phi(self, phi, 0, error);
    }
}

impl<S: Space, P: Projector<S>> QFunction<S> for Linear<S, P> {
    fn evaluate_action(&self, input: &S::Repr, action: usize) -> f64 {
        let phi = self.projector.project(input);

        self.evaluate_action_phi(&phi, action)
    }

    fn update_action(&mut self, input: &S::Repr, action: usize, error: f64) {
        let phi = self.projector.project(input);

        self.update_action_phi(&phi, action, error);
    }

    fn evaluate_phi(&self, phi: &Projection) -> Vec<f64> {
        match phi {
            &Projection::Dense(ref dense_phi) => (self.weights.t().dot(&(dense_phi/phi.z()))).into_raw_vec(),
            &Projection::Sparse(ref sparse_phi) =>
                (0..self.weights.cols()).map(|c| {
                    sparse_phi.iter().fold(0.0, |acc, idx| acc + self.weights[(*idx, c)])
                }).collect(),
        }
    }

    fn evaluate_action_phi(&self, phi: &Projection, action: usize) -> f64 {
        let col = self.weights.column(action);

        match phi {
            &Projection::Dense(ref dense_phi) => col.dot(&(dense_phi/phi.z())),
            &Projection::Sparse(ref sparse_phi) =>
                sparse_phi.iter().fold(0.0, |acc, idx| acc + col[*idx]),
        }
    }

    fn update_phi(&mut self, phi: &Projection, errors: Vec<f64>) {
        let z = phi.z();
        let sf = 1.0/z;

        match phi {
            &Projection::Dense(ref dense_phi) => {
                let view = dense_phi.view().into_shape((self.weights.rows(), 1)).unwrap();
                let error_matrix = ArrayView::from_shape((1, self.weights.cols()), errors.as_slice())
                    .unwrap();

                self.weights.scaled_add(sf, &view.dot(&error_matrix))
            },
            &Projection::Sparse(ref sparse_phi) => {
                for c in 0..self.weights.cols() {
                    let mut col = self.weights.column_mut(c);
                    let error = errors[c];

                    for idx in sparse_phi {
                        col[*idx] += error
                    }
                }
            },
        }
    }

    fn update_action_phi(&mut self, phi: &Projection, action: usize, error: f64) {
        let mut col = self.weights.column_mut(action);
        let scaled_error = error/phi.z();

        match phi {
            &Projection::Dense(ref dense_phi) => col.scaled_add(scaled_error, dense_phi),
            &Projection::Sparse(ref sparse_phi) => {
                for idx in sparse_phi {
                    col[*idx] += scaled_error
                }
            },
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate seahash;

    use super::*;
    use std::hash::BuildHasherDefault;
    use fa::projection::TileCoding;

    type SHBuilder = BuildHasherDefault<seahash::SeaHasher>;

    #[test]
    fn test_dense_update_eval() {
        let p = TileCoding::new(SHBuilder::default(), 4, 100);
        let mut f = Linear::new(p.clone(), 1);

        let input = vec![5.0];

        f.update(&input, 50.0);
        let out: f64 = f.evaluate(&input);

        assert!(out > 0.0);
    }

    #[test]
    fn test_dense_eval_phi() {
        let p = TileCoding::new(SHBuilder::default(), 4, 100);
        let f = Linear::new(p.clone(), 1);

        let input = vec![5.0];

        let out: f64 = f.evaluate(&input);
        let out_alt_v: f64 = VFunction::evaluate_phi(&f, &p.project(&input));
        let out_alt_q: Vec<f64> = QFunction::evaluate_phi(&f, &p.project(&input));

        assert_eq!(out, out_alt_v);
        assert_eq!(vec![out], out_alt_q);
    }
}
<|endoftext|>


<commit_before>pub use self::os::{flock};
pub use self::os::{F_DUPFD};
pub use self::os::{F_DUPFD_CLOEXEC};
pub use self::os::{F_GETFD};
pub use self::os::{F_SETFD};
pub use self::os::{F_GETFL};
pub use self::os::{F_SETFL};
pub use self::os::{F_GETLK};
pub use self::os::{F_SETLK};
pub use self::os::{F_SETLKW};
pub use self::os::{F_GETOWN};
pub use self::os::{F_SETOWN};
pub use self::os::{FD_CLOEXEC};
pub use self::os::{F_RDLCK};
pub use self::os::{F_UNLCK};
pub use self::os::{F_WRLCK};
pub use self::os::{O_CLOEXEC};
pub use self::os::{O_CREAT};
pub use self::os::{O_DIRECTORY};
pub use self::os::{O_EXCL};
pub use self::os::{O_NOCTTY};
pub use self::os::{O_NOFOLLOW};
pub use self::os::{O_TRUNC};
pub use self::os::{O_APPEND};
pub use self::os::{O_DSYNC};
pub use self::os::{O_NONBLOCK};
pub use self::os::{O_RSYNC};
pub use self::os::{O_SYNC};
pub use self::os::{O_ACCMODE};
pub use self::os::{O_RDONLY};
pub use self::os::{O_RDWR};
pub use self::os::{O_WRONLY};
pub use self::os::{AT_FDCWD};
pub use self::os::{AT_EACCESS};
pub use self::os::{AT_SYMLINK_NOFOLLOW};
pub use self::os::{AT_SYMLINK_FOLLOW};
pub use self::os::{AT_REMOVEDIR};
pub use self::os::{POSIX_FADV_DONTNEED};
pub use self::os::{POSIX_FADV_NOREUSE};
pub use self::os::{POSIX_FADV_NORMAL};
pub use self::os::{POSIX_FADV_RANDOM};
pub use self::os::{POSIX_FADV_SEQUENTIAL};
pub use self::os::{POSIX_FADV_WILLNEED};

#[cfg(target_os = "linux")]
#[path = "linux/mod.rs"]
mod os;

pub fn creat<T: ::NTStr>(file: &T, mode: ::sys::types::mode_t) -> ::int_t {
    extern { fn creat(file: *const ::char_t, mode: ::sys::types::mode_t) -> ::int_t; }
    unsafe { creat(file.as_ptr(), mode) }
}

pub fn posix_fadvise(fd: ::int_t, off: ::sys::types::off_t, len: ::sys::types::off_t,
                     advise: ::int_t) -> ::int_t {
    extern {
        fn posix_fadvise(fd: ::int_t, off: ::sys::types::off_t, len: ::sys::types::off_t,
                         advise: ::int_t) -> ::int_t;
    }
    unsafe { posix_fadvise(fd, off, len, advise) }
}

pub fn posix_fallocate(fd: ::int_t, off: ::sys::types::off_t,
                       len: ::sys::types::off_t) -> ::int_t {
    extern {  fn posix_fallocate(fd: ::int_t, off: ::sys::types::off_t,
                                 len: ::sys::types::off_t) -> ::int_t; }
    unsafe { posix_fallocate(fd, off, len) }
}

extern "C" {
    pub fn fcntl(fd: ::int_t, cmd: ::int_t, ...) -> ::int_t;
    pub fn open(file: *const ::char_t, oflag: ::int_t, ...) -> ::int_t;
    pub fn openat(fd: ::int_t, file: *const ::char_t, oflag: ::int_t, ...) -> ::int_t;
}
<commit_msg>fcntl.h<commit_after>pub use self::os::{flock};
pub use self::os::{F_DUPFD};
pub use self::os::{F_DUPFD_CLOEXEC};
pub use self::os::{F_GETFD};
pub use self::os::{F_SETFD};
pub use self::os::{F_GETFL};
pub use self::os::{F_SETFL};
pub use self::os::{F_GETLK};
pub use self::os::{F_SETLK};
pub use self::os::{F_SETLKW};
pub use self::os::{F_GETOWN};
pub use self::os::{F_SETOWN};
pub use self::os::{FD_CLOEXEC};
pub use self::os::{F_RDLCK};
pub use self::os::{F_UNLCK};
pub use self::os::{F_WRLCK};
pub use self::os::{O_CLOEXEC};
pub use self::os::{O_CREAT};
pub use self::os::{O_DIRECTORY};
pub use self::os::{O_EXCL};
pub use self::os::{O_NOCTTY};
pub use self::os::{O_NOFOLLOW};
pub use self::os::{O_TRUNC};
pub use self::os::{O_APPEND};
pub use self::os::{O_DSYNC};
pub use self::os::{O_NONBLOCK};
pub use self::os::{O_RSYNC};
pub use self::os::{O_SYNC};
pub use self::os::{O_ACCMODE};
pub use self::os::{O_RDONLY};
pub use self::os::{O_RDWR};
pub use self::os::{O_WRONLY};
pub use self::os::{AT_FDCWD};
pub use self::os::{AT_EACCESS};
pub use self::os::{AT_SYMLINK_NOFOLLOW};
pub use self::os::{AT_SYMLINK_FOLLOW};
pub use self::os::{AT_REMOVEDIR};
pub use self::os::{POSIX_FADV_DONTNEED};
pub use self::os::{POSIX_FADV_NOREUSE};
pub use self::os::{POSIX_FADV_NORMAL};
pub use self::os::{POSIX_FADV_RANDOM};
pub use self::os::{POSIX_FADV_SEQUENTIAL};
pub use self::os::{POSIX_FADV_WILLNEED};

use {NTStr, int_t, char_t};
use sys::types::{mode_t, off_t};

#[cfg(target_os = "linux")]
#[path = "linux/mod.rs"]
mod os;

pub fn creat<T: NTStr>(file: &T, mode: mode_t) -> int_t {
    extern { fn creat(file: *const char_t, mode: mode_t) -> int_t; }
    unsafe { creat(file.as_ptr(), mode) }
}

pub fn posix_fadvise(fd: int_t, off: off_t, len: off_t, advise: int_t) -> int_t {
    extern {
        fn posix_fadvise(fd: int_t, off: off_t, len: off_t, advise: int_t) -> int_t;
    }
    unsafe { posix_fadvise(fd, off, len, advise) }
}

pub fn posix_fallocate(fd: int_t, off: off_t, len: off_t) -> int_t {
    extern { fn posix_fallocate(fd: int_t, off: off_t, len: off_t) -> int_t; }
    unsafe { posix_fallocate(fd, off, len) }
}

extern "C" {
    pub fn fcntl(fd: int_t, cmd: int_t, ...) -> int_t;
    pub fn open(file: *const char_t, oflag: int_t, ...) -> int_t;
    pub fn openat(fd: int_t, file: *const char_t, oflag: int_t, ...) -> int_t;
}
<|endoftext|>


<commit_before>extern crate hyper;

use std::io::Read;
use std::iter::Iterator;

use hyper::Client;
use hyper::net::HttpsConnector;
use hyper::Error;
use hyper_native_tls::NativeTlsClient;

/// Fetches a gitignore file from [gitignore.io](https://gitignore.io)
///
/// # Example
///
/// ```
/// mod gitignore
///
/// let gi = gitignore::get_gitignore(&vec!["rust".to_string, "vim".to_string()]).unwrap();
/// ```
pub fn get_gitignore(targets: &Vec<String>) -> Result<String, hyper::Error> {
    let mut request = "https://www.gitignore.io/api/".to_string();
    for (i, target) in targets.iter().enumerate() {
        request.push_str(&target);

        // skip last ","
        if i != targets.len() - 1 {
            request.push_str(",");
        }
    }

    let ssl = NativeTlsClient::new().unwrap();
    let con = HttpsConnector::new(ssl);
    let client = Client::with_connector(con);
    let url = request.parse::<hyper::Url>().unwrap();
    let req = client.get(url);
    let mut resp = match req.send() {
        Ok(r) => r,
        Err(err) => return Err(err),
    };

    let mut body = String::new();
    match resp.read_to_string(&mut body) {
        Ok(_) => {}
        Err(e) => return Err(Error::from(e)),
    };
    Ok(body)
}

#[test]
fn test_get_gitignore() {
    let expected = "
# Created by https://www.gitignore.io/api/rust

### Rust ###
# Generated by Cargo
# will have compiled files and executables
/target/

# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
# More information here http://doc.crates.io/guide.html#cargotoml-vs-cargolock
Cargo.lock

# These are backup files generated by rustfmt
**/*.rs.bk

# End of https://www.gitignore.io/api/rust
";
    let actual = get_gitignore(&vec!["rust".to_string()]).unwrap();
    assert_eq!(expected.to_string(), actual);
}
<commit_msg>Make error handling more readable<commit_after>extern crate hyper;

use std::io::Read;
use std::iter::Iterator;

use hyper::Client;
use hyper::net::HttpsConnector;
use hyper_native_tls::NativeTlsClient;

/// Fetches a gitignore file from [gitignore.io](https://gitignore.io)
///
/// # Example
///
/// ```
/// mod gitignore
///
/// let gi = gitignore::get_gitignore(&vec!["rust".to_string, "vim".to_string()]).unwrap();
/// ```
pub fn get_gitignore(targets: &Vec<String>) -> Result<String, hyper::Error> {
    let mut request = "https://www.gitignore.io/api/".to_string();
    for (i, target) in targets.iter().enumerate() {
        request.push_str(&target);

        // skip last ","
        if i != targets.len() - 1 {
            request.push_str(",");
        }
    }

    let ssl = NativeTlsClient::new().unwrap();
    let con = HttpsConnector::new(ssl);
    let client = Client::with_connector(con);
    let url = request.parse::<hyper::Url>().unwrap();
    let req = client.get(url);
    let mut resp = match req.send() {
        Ok(r) => r,
        Err(err) => return Err(err),
    };

    let mut body = String::new();
    resp.read_to_string(&mut body)?;

    Ok(body)
}

#[test]
fn test_get_gitignore() {
    let expected = "
# Created by https://www.gitignore.io/api/rust

### Rust ###
# Generated by Cargo
# will have compiled files and executables
/target/

# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
# More information here http://doc.crates.io/guide.html#cargotoml-vs-cargolock
Cargo.lock

# These are backup files generated by rustfmt
**/*.rs.bk

# End of https://www.gitignore.io/api/rust
";
    let actual = get_gitignore(&vec!["rust".to_string()]).unwrap();
    assert_eq!(expected.to_string(), actual);
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Utils for forming gzchunked streams.

/// Encodes the given iterator over protobuf messages into the gzchunked format.
///
/// This is a streaming encoder that lazily encodes the data and can be used to
/// effectively process megabytes of data. The function returns an iterator that
/// yields parts of the gzchunked file.
///
/// This variant uses default encoding settings. To customize them, one can use
/// the [`encode_with_opts`] function.
///
/// [`encode_with_opts`]: fn.encode_with_opts.html
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
/// use std::io::Write as _;
///
/// let items = vec! {
///     String::from("foo"),
///     String::from("bar"),
///     String::from("baz"),
/// };
///
/// let chunks = rrg::gzchunked::encode(items.into_iter());
/// for (idx, chunk) in chunks.enumerate() {
///     let mut file = File::open(format!("output.gzc.{}", idx)).unwrap();
///     file.write_all(chunk.unwrap().as_slice());
/// }
/// ```
pub fn encode<I, M>(iter: I) -> Encode<I>
where
    I: Iterator<Item=M>,
    M: prost::Message,
{
    encode_with_opts(iter, EncodeOpts::default())
}

/// Encodes the given iterator over protobuf messages into the gzchunked format.
///
/// This is a variant of the [`encode`] function that allows customization of
/// encoding parameters. Refer to its documentation for more details.
///
/// [`encode`]: fn.encode.html
pub fn encode_with_opts<I, M>(iter: I, opts: EncodeOpts) -> Encode<I>
where
    I: Iterator<Item=M>,
    M: prost::Message,
{
    Encode::with_opts(iter, opts)
}

/// Decodes an iterator over gzchunked file parts into a stream of messages.
///
/// This is a streaming decoder that performs the decoding in a lazy way and can
/// be used to effectively process megabytes of data.
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
///
/// let paths = ["foo.gzc.1", "foo.gzc.2", "foo.gzc.3"];
/// let files = paths.iter().map(|path| File::open(path).unwrap());
///
/// for (idx, msg) in rrg::gzchunked::decode(files).enumerate() {
///     let msg: String = msg.unwrap();
///     println!("item #{}: {:?}", idx, msg);
/// }
/// ```
pub fn decode<I, R, M>(iter: I) -> impl Iterator<Item=std::io::Result<M>>
where
    I: Iterator<Item=R>,
    R: std::io::Read,
    M: prost::Message + Default,
{
    let parts = iter.map(flate2::read::GzDecoder::new);
    crate::chunked::decode(crate::io::IterReader::new(parts))
}

/// A wrapper type for gzip compression level.
#[derive(Clone, Copy, Debug)]
pub struct Compression(flate2::Compression);

/// A type for defining gzip compression level for gzchunked.
impl Compression {

    /// Creates a new compression descriptor at the specified level.
    ///
    /// The lower the number the worse the compression (with 0 meaning no
    /// compression at all).
    pub fn new(level: u32) -> Compression {
        Compression(flate2::Compression::new(level))
    }

    /// Creates a new compression descriptor with disabled compression.
    pub fn none() -> Compression {
        Compression(flate2::Compression::none())
    }

    /// Creates a new compression descriptor with highest compression.
    pub fn best() -> Compression {
        Compression(flate2::Compression::best())
    }
}

impl Default for Compression {

    fn default() -> Compression {
        Compression(flate2::Compression::new(5))
    }
}

/// Options and flags that configure encoding into the gzchuned format.
#[derive(Clone, Copy, Debug)]
pub struct EncodeOpts {
    /// Compression level used for the gzip encoding.
    pub compression: Compression,
    /// A rough file size limit for parts of the output file.
    pub part_size: u64,
}

impl Default for EncodeOpts {

    fn default() -> EncodeOpts {
        EncodeOpts {
            compression: Compression::default(),
            part_size: 1 * 1024 * 1024, // 1 MiB.
        }
    }
}

/// Streaming encoder for the gzchunked format.
///
/// It implements the `Iterator` trait, lazily polling the underlying iterator
/// over Protocol Buffers messages as more parts are needed.
///
/// Instances of this type can be constructed using the [`encode_with_opts`] or
/// [`encode`] function.
///
/// [`encode`]: fn.encode.html
/// [`encode_with_opts`]: fn.encode_with_opts.html
pub struct Encode<I> {
    chunked: crate::chunked::Encode<I>,
    opts: EncodeOpts,
}

impl<I> Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    /// Creates a new encoder instance with the specified options.
    fn with_opts(iter: I, opts: EncodeOpts) -> Encode<I> {
        Encode {
            chunked: crate::chunked::encode(iter),
            opts: opts,
        }
    }

    /// Obtains the next part of the output file (if available).
    fn next_part(&mut self) -> std::io::Result<Option<Vec<u8>>> {
        use crate::io::copy_until;

        let compression = self.opts.compression.0;
        let part_size = self.opts.part_size;

        let mut encoder = flate2::write::GzEncoder::new(vec!(), compression);
        let len = copy_until(&mut self.chunked, &mut encoder, |_, encoder| {
            encoder.get_ref().len() as u64 >= part_size
        })?;

        if len == 0 {
            Ok(None)
        } else {
            Ok(Some(encoder.finish()?))
        }
    }
}

impl<I> Iterator for Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    type Item = std::io::Result<Vec<u8>>;

    fn next(&mut self) -> Option<std::io::Result<Vec<u8>>> {
        self.next_part().transpose()
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn test_encode_with_empty_iter() {
        let mut iter = encode(std::iter::empty::<()>())
            .map(Result::unwrap);

        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_decode_with_empty_iter() {
        let mut iter = decode::<_, _, ()>(std::iter::empty::<&[u8]>())
            .map(Result::unwrap);

        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_single_item_iter() {
        let chunks = encode(std::iter::once(String::from("foo")))
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(String::from("foo")));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_multiple_items_iter() {
        let data = vec! {
            String::from("foo"),
            String::from("bar"),
            String::from("baz"),
        };

        let chunks = encode(data.into_iter())
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(String::from("foo")));
        assert_eq!(iter.next(), Some(String::from("bar")));
        assert_eq!(iter.next(), Some(String::from("baz")));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_empty_items_iter() {
        let data = vec!((), (), ());

        let chunks = encode(data.into_iter())
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_many_items_iter() {
        let sample = rand::random::<[u8; 32]>().to_vec();
        let items = std::iter::repeat(sample.clone()).take(32 * 1024);

        let opts = EncodeOpts {
            compression: Compression::default(),
            part_size: 4 * 1024,
        };

        let chunks = encode_with_opts(items, opts)
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode::<_, _, Vec<u8>>(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert!(iter.all(|item| item == sample));
    }

    #[test]
    fn test_encode_and_decode_with_no_compression() {
        let sample = rand::random::<[u8; 32]>().to_vec();
        let items = std::iter::repeat(sample.clone()).take(32 * 1024);

        let opts = EncodeOpts {
            compression: Compression::none(),
            part_size: 4 * 1024,
        };

        let chunks = encode_with_opts(items, opts)
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode::<_, _, Vec<u8>>(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert!(iter.all(|item| item == sample));
    }

    #[test]
    fn test_encode_and_decode_with_best_compression() {
        let sample = rand::random::<[u8; 32]>().to_vec();
        let items = std::iter::repeat(sample.clone()).take(32 * 1024);

        let opts = EncodeOpts {
            compression: Compression::best(),
            part_size: 4 * 1024,
        };

        let chunks = encode_with_opts(items, opts)
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode::<_, _, Vec<u8>>(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert!(iter.all(|item| item == sample));
    }
}
<commit_msg>Refactor type signatures for gzchunked utilities.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Utils for forming gzchunked streams.

/// Encodes the given iterator over protobuf messages into the gzchunked format.
///
/// This is a streaming encoder that lazily encodes the data and can be used to
/// effectively process megabytes of data. The function returns an iterator that
/// yields parts of the gzchunked file.
///
/// This variant uses default encoding settings. To customize them, one can use
/// the [`encode_with_opts`] function.
///
/// [`encode_with_opts`]: fn.encode_with_opts.html
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
/// use std::io::Write as _;
///
/// let items = vec! {
///     String::from("foo"),
///     String::from("bar"),
///     String::from("baz"),
/// };
///
/// let chunks = rrg::gzchunked::encode(items.into_iter());
/// for (idx, chunk) in chunks.enumerate() {
///     let mut file = File::open(format!("output.gzc.{}", idx)).unwrap();
///     file.write_all(chunk.unwrap().as_slice());
/// }
/// ```
pub fn encode<I>(iter: I) -> Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    encode_with_opts(iter, EncodeOpts::default())
}

/// Encodes the given iterator over protobuf messages into the gzchunked format.
///
/// This is a variant of the [`encode`] function that allows customization of
/// encoding parameters. Refer to its documentation for more details.
///
/// [`encode`]: fn.encode.html
pub fn encode_with_opts<I>(iter: I, opts: EncodeOpts) -> Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    Encode::with_opts(iter, opts)
}

/// Decodes an iterator over gzchunked file parts into a stream of messages.
///
/// This is a streaming decoder that performs the decoding in a lazy way and can
/// be used to effectively process megabytes of data.
///
/// # Examples
///
/// ```no_run
/// use std::fs::File;
///
/// let paths = ["foo.gzc.1", "foo.gzc.2", "foo.gzc.3"];
/// let files = paths.iter().map(|path| File::open(path).unwrap());
///
/// for (idx, msg) in rrg::gzchunked::decode(files).enumerate() {
///     let msg: String = msg.unwrap();
///     println!("item #{}: {:?}", idx, msg);
/// }
/// ```
pub fn decode<I, M>(iter: I) -> impl Iterator<Item=std::io::Result<M>>
where
    I: Iterator,
    I::Item: std::io::Read,
    M: prost::Message + Default,
{
    let parts = iter.map(flate2::read::GzDecoder::new);
    crate::chunked::decode(crate::io::IterReader::new(parts))
}

/// A wrapper type for gzip compression level.
#[derive(Clone, Copy, Debug)]
pub struct Compression(flate2::Compression);

/// A type for defining gzip compression level for gzchunked.
impl Compression {

    /// Creates a new compression descriptor at the specified level.
    ///
    /// The lower the number the worse the compression (with 0 meaning no
    /// compression at all).
    pub fn new(level: u32) -> Compression {
        Compression(flate2::Compression::new(level))
    }

    /// Creates a new compression descriptor with disabled compression.
    pub fn none() -> Compression {
        Compression(flate2::Compression::none())
    }

    /// Creates a new compression descriptor with highest compression.
    pub fn best() -> Compression {
        Compression(flate2::Compression::best())
    }
}

impl Default for Compression {

    fn default() -> Compression {
        Compression(flate2::Compression::new(5))
    }
}

/// Options and flags that configure encoding into the gzchuned format.
#[derive(Clone, Copy, Debug)]
pub struct EncodeOpts {
    /// Compression level used for the gzip encoding.
    pub compression: Compression,
    /// A rough file size limit for parts of the output file.
    pub part_size: u64,
}

impl Default for EncodeOpts {

    fn default() -> EncodeOpts {
        EncodeOpts {
            compression: Compression::default(),
            part_size: 1 * 1024 * 1024, // 1 MiB.
        }
    }
}

/// Streaming encoder for the gzchunked format.
///
/// It implements the `Iterator` trait, lazily polling the underlying iterator
/// over Protocol Buffers messages as more parts are needed.
///
/// Instances of this type can be constructed using the [`encode_with_opts`] or
/// [`encode`] function.
///
/// [`encode`]: fn.encode.html
/// [`encode_with_opts`]: fn.encode_with_opts.html
pub struct Encode<I> {
    chunked: crate::chunked::Encode<I>,
    opts: EncodeOpts,
}

impl<I> Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    /// Creates a new encoder instance with the specified options.
    fn with_opts(iter: I, opts: EncodeOpts) -> Encode<I> {
        Encode {
            chunked: crate::chunked::encode(iter),
            opts: opts,
        }
    }

    /// Obtains the next part of the output file (if available).
    fn next_part(&mut self) -> std::io::Result<Option<Vec<u8>>> {
        use crate::io::copy_until;

        let compression = self.opts.compression.0;
        let part_size = self.opts.part_size;

        let mut encoder = flate2::write::GzEncoder::new(vec!(), compression);
        let len = copy_until(&mut self.chunked, &mut encoder, |_, encoder| {
            encoder.get_ref().len() as u64 >= part_size
        })?;

        if len == 0 {
            Ok(None)
        } else {
            Ok(Some(encoder.finish()?))
        }
    }
}

impl<I> Iterator for Encode<I>
where
    I: Iterator,
    I::Item: prost::Message,
{
    type Item = std::io::Result<Vec<u8>>;

    fn next(&mut self) -> Option<std::io::Result<Vec<u8>>> {
        self.next_part().transpose()
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn test_encode_with_empty_iter() {
        let mut iter = encode(std::iter::empty::<()>())
            .map(Result::unwrap);

        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_decode_with_empty_iter() {
        let mut iter = decode::<_, ()>(std::iter::empty::<&[u8]>())
            .map(Result::unwrap);

        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_single_item_iter() {
        let chunks = encode(std::iter::once(String::from("foo")))
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(String::from("foo")));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_multiple_items_iter() {
        let data = vec! {
            String::from("foo"),
            String::from("bar"),
            String::from("baz"),
        };

        let chunks = encode(data.into_iter())
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(String::from("foo")));
        assert_eq!(iter.next(), Some(String::from("bar")));
        assert_eq!(iter.next(), Some(String::from("baz")));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_empty_items_iter() {
        let data = vec!((), (), ());

        let chunks = encode(data.into_iter())
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), Some(b"".to_vec()));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_encode_and_decode_with_many_items_iter() {
        let sample = rand::random::<[u8; 32]>().to_vec();
        let items = std::iter::repeat(sample.clone()).take(32 * 1024);

        let opts = EncodeOpts {
            compression: Compression::default(),
            part_size: 4 * 1024,
        };

        let chunks = encode_with_opts(items, opts)
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode::<_, Vec<u8>>(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert!(iter.all(|item| item == sample));
    }

    #[test]
    fn test_encode_and_decode_with_no_compression() {
        let sample = rand::random::<[u8; 32]>().to_vec();
        let items = std::iter::repeat(sample.clone()).take(32 * 1024);

        let opts = EncodeOpts {
            compression: Compression::none(),
            part_size: 4 * 1024,
        };

        let chunks = encode_with_opts(items, opts)
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode::<_, Vec<u8>>(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert!(iter.all(|item| item == sample));
    }

    #[test]
    fn test_encode_and_decode_with_best_compression() {
        let sample = rand::random::<[u8; 32]>().to_vec();
        let items = std::iter::repeat(sample.clone()).take(32 * 1024);

        let opts = EncodeOpts {
            compression: Compression::best(),
            part_size: 4 * 1024,
        };

        let chunks = encode_with_opts(items, opts)
            .map(Result::unwrap)
            .collect::<Vec<_>>();

        let mut iter = decode::<_, Vec<u8>>(chunks.iter().map(Vec::as_slice))
            .map(Result::unwrap);

        assert!(iter.all(|item| item == sample));
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;
use std::io::{Reader, Writer};

pub fn read_at_least<R : Reader>(reader : &mut R,
                                 buf: &mut [u8],
                                 min_bytes : uint) -> uint {
    let mut pos = 0;
    let bufLen = buf.len();
    while pos < min_bytes {
        let buf1 = buf.mut_slice(pos, bufLen);
        match reader.read(buf1) {
            None => fail!("premature EOF?"),
            Some(n) => pos += n
        }
    }
    return pos;
}

pub struct BufferedInputStream<'a, R> {
    priv inner : &'a mut R,
    priv buf : ~[u8],
    priv pos : uint,
    priv cap : uint
}

impl<'a, R: Reader> BufferedInputStream<'a, R> {
    pub fn new<'a> (r : &'a mut R) -> BufferedInputStream<'a, R> {
        let mut result = BufferedInputStream {
            inner : r,
            buf : std::vec::with_capacity(8192),
            pos : 0,
            cap : 0
        };
        unsafe {
            std::vec::raw::set_len(&mut result.buf, 8192)
        }
        return result;
    }

    pub fn skip(&mut self, mut bytes : uint) {
        let available = self.cap - self.pos;
        if bytes <= available {
            self.pos += bytes;
        } else {
            bytes -= available;
            if (bytes <= self.buf.len()) {
                //# Read the next buffer-full.
                let n = read_at_least(self.inner, self.buf, bytes);
                self.pos = bytes;
                self.cap = n;

            } else {
                //# Forward large skip to the underlying stream.
                fail!("TODO")
            }
        }
    }

    pub fn getReadBuffer(&self) -> (*u8, *u8) {
        unsafe {
            (self.buf.unsafe_ref(self.pos),
             self.buf.unsafe_ref(self.cap))
        }
    }
}

impl<'a, R: Reader> Reader for BufferedInputStream<'a, R> {
    fn eof(&mut self) -> bool {
        self.inner.eof()
    }

    fn read(&mut self, dst: &mut [u8]) -> Option<uint> {
        let mut num_bytes = dst.len();
        if (num_bytes <= self.cap - self.pos) {
            //# Serve from the current buffer.
            std::vec::bytes::copy_memory(dst,
                                         self.buf.slice(self.pos, self.cap),
                                         num_bytes);
            self.pos += num_bytes;
            return Some(num_bytes);
        } else {
            //# Copy current available into destination.

            std::vec::bytes::copy_memory(dst,
                                         self.buf.slice(self.pos, self.cap),
                                         self.cap - self.pos);
            let fromFirstBuffer = self.cap - self.pos;

            let dst1 = dst.mut_slice(fromFirstBuffer, num_bytes);
            num_bytes -= fromFirstBuffer;
            if (num_bytes <= self.buf.len()) {
                //# Read the next buffer-full.
                let n = read_at_least(self.inner, self.buf, num_bytes);
                std::vec::bytes::copy_memory(dst1,
                                             self.buf.slice(0, num_bytes),
                                             num_bytes);
                self.cap = n;
                self.pos = num_bytes;
                return Some(fromFirstBuffer + num_bytes);
            } else {
                //# Forward large read to the underlying stream.
                self.pos = 0;
                self.cap = 0;
                return Some(fromFirstBuffer + read_at_least(self.inner, dst1, num_bytes));
            }
        }
    }
}

pub struct BufferedOutputStream<'a, W> {
    priv inner: &'a mut W,
    priv buf: ~[u8],
    priv pos: uint
}

impl<'self, W: Writer> BufferedOutputStream<'self, W> {

    pub fn new<'a> (w : &'a mut W) -> BufferedOutputStream<'a, W> {
        let mut result = BufferedOutputStream {
            inner: w,
            buf : std::vec::with_capacity(8192),
            pos : 0
        };
        unsafe {
            std::vec::raw::set_len(&mut result.buf, 8192);
        }
        return result;
    }

    #[inline]
    pub fn getWriteBuffer(&mut self) -> (*mut u8, *mut u8) {
        let len = self.buf.len();
        unsafe {
            (self.buf.unsafe_mut_ref(self.pos),
             self.buf.unsafe_mut_ref(len))
        }
    }

    #[inline]
    pub fn write_ptr(&mut self, ptr: *mut u8, size: uint) {
        unsafe {
            let easyCase = ptr == self.buf.unsafe_mut_ref(self.pos);
            if easyCase {
                self.pos += size;
            } else {
                do std::vec::raw::mut_buf_as_slice::<u8,()>(ptr, size) |buf| {
                    self.write(buf);
                }
            }
        }
    }
}

impl<'self, W: Writer> Writer for BufferedOutputStream<'self, W> {
    #[inline]
    fn write(&mut self, buf: &[u8]) {
        let available = self.buf.len() - self.pos;
        let mut size = buf.len();
        if size <= available {
            let dst = self.buf.mut_slice_from(self.pos);
            std::vec::bytes::copy_memory(dst, buf, buf.len());
            self.pos += size;
        } else if size <= self.buf.len() {
            //# Too much for this buffer, but not a full buffer's
            //# worth, so we'll go ahead and copy.
            {
                let dst = self.buf.mut_slice_from(self.pos);
                std::vec::bytes::copy_memory(dst, buf, available);
            }
            self.inner.write(self.buf);

            size -= available;
            let src = buf.slice(available, buf.len());
            let dst = self.buf.mut_slice_from(0);
            std::vec::bytes::copy_memory(dst, src, size);
            self.pos = size;
        } else {
            //# Writing so much data that we might as well write
            //# directly to avoid a copy.
            self.inner.write(self.buf.slice(0, self.pos));
            self.pos = 0;
            self.inner.write(buf);
        }
    }

    fn flush(&mut self) {
        if (self.pos > 0) {
            self.inner.write(self.buf.slice(0, self.pos));
            self.pos = 0;
        }
    }
}
<commit_msg>move unsafety into function sigs<commit_after>/*
 * Copyright (c) 2013, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;
use std::io::{Reader, Writer};

pub fn read_at_least<R : Reader>(reader : &mut R,
                                 buf: &mut [u8],
                                 min_bytes : uint) -> uint {
    let mut pos = 0;
    let bufLen = buf.len();
    while pos < min_bytes {
        let buf1 = buf.mut_slice(pos, bufLen);
        match reader.read(buf1) {
            None => fail!("premature EOF?"),
            Some(n) => pos += n
        }
    }
    return pos;
}

pub struct BufferedInputStream<'a, R> {
    priv inner : &'a mut R,
    priv buf : ~[u8],
    priv pos : uint,
    priv cap : uint
}

impl<'a, R: Reader> BufferedInputStream<'a, R> {
    pub fn new<'a> (r : &'a mut R) -> BufferedInputStream<'a, R> {
        let mut result = BufferedInputStream {
            inner : r,
            buf : std::vec::with_capacity(8192),
            pos : 0,
            cap : 0
        };
        unsafe {
            std::vec::raw::set_len(&mut result.buf, 8192)
        }
        return result;
    }

    pub fn skip(&mut self, mut bytes : uint) {
        let available = self.cap - self.pos;
        if bytes <= available {
            self.pos += bytes;
        } else {
            bytes -= available;
            if (bytes <= self.buf.len()) {
                //# Read the next buffer-full.
                let n = read_at_least(self.inner, self.buf, bytes);
                self.pos = bytes;
                self.cap = n;

            } else {
                //# Forward large skip to the underlying stream.
                fail!("TODO")
            }
        }
    }

    pub unsafe fn getReadBuffer(&self) -> (*u8, *u8) {
        (self.buf.unsafe_ref(self.pos), self.buf.unsafe_ref(self.cap))
    }
}

impl<'a, R: Reader> Reader for BufferedInputStream<'a, R> {
    fn eof(&mut self) -> bool {
        self.inner.eof()
    }

    fn read(&mut self, dst: &mut [u8]) -> Option<uint> {
        let mut num_bytes = dst.len();
        if (num_bytes <= self.cap - self.pos) {
            //# Serve from the current buffer.
            std::vec::bytes::copy_memory(dst,
                                         self.buf.slice(self.pos, self.cap),
                                         num_bytes);
            self.pos += num_bytes;
            return Some(num_bytes);
        } else {
            //# Copy current available into destination.

            std::vec::bytes::copy_memory(dst,
                                         self.buf.slice(self.pos, self.cap),
                                         self.cap - self.pos);
            let fromFirstBuffer = self.cap - self.pos;

            let dst1 = dst.mut_slice(fromFirstBuffer, num_bytes);
            num_bytes -= fromFirstBuffer;
            if (num_bytes <= self.buf.len()) {
                //# Read the next buffer-full.
                let n = read_at_least(self.inner, self.buf, num_bytes);
                std::vec::bytes::copy_memory(dst1,
                                             self.buf.slice(0, num_bytes),
                                             num_bytes);
                self.cap = n;
                self.pos = num_bytes;
                return Some(fromFirstBuffer + num_bytes);
            } else {
                //# Forward large read to the underlying stream.
                self.pos = 0;
                self.cap = 0;
                return Some(fromFirstBuffer + read_at_least(self.inner, dst1, num_bytes));
            }
        }
    }
}

pub struct BufferedOutputStream<'a, W> {
    priv inner: &'a mut W,
    priv buf: ~[u8],
    priv pos: uint
}

impl<'self, W: Writer> BufferedOutputStream<'self, W> {

    pub fn new<'a> (w : &'a mut W) -> BufferedOutputStream<'a, W> {
        let mut result = BufferedOutputStream {
            inner: w,
            buf : std::vec::with_capacity(8192),
            pos : 0
        };
        unsafe {
            std::vec::raw::set_len(&mut result.buf, 8192);
        }
        return result;
    }

    #[inline]
    pub unsafe fn getWriteBuffer(&mut self) -> (*mut u8, *mut u8) {
        let len = self.buf.len();
        (self.buf.unsafe_mut_ref(self.pos), self.buf.unsafe_mut_ref(len))
    }

    #[inline]
    pub unsafe fn write_ptr(&mut self, ptr: *mut u8, size: uint) {
        let easyCase = ptr == self.buf.unsafe_mut_ref(self.pos);
        if easyCase {
            self.pos += size;
        } else {
            do std::vec::raw::mut_buf_as_slice::<u8,()>(ptr, size) |buf| {
                self.write(buf);
            }
        }
    }
}

impl<'self, W: Writer> Writer for BufferedOutputStream<'self, W> {
    #[inline]
    fn write(&mut self, buf: &[u8]) {
        let available = self.buf.len() - self.pos;
        let mut size = buf.len();
        if size <= available {
            let dst = self.buf.mut_slice_from(self.pos);
            std::vec::bytes::copy_memory(dst, buf, buf.len());
            self.pos += size;
        } else if size <= self.buf.len() {
            //# Too much for this buffer, but not a full buffer's
            //# worth, so we'll go ahead and copy.
            {
                let dst = self.buf.mut_slice_from(self.pos);
                std::vec::bytes::copy_memory(dst, buf, available);
            }
            self.inner.write(self.buf);

            size -= available;
            let src = buf.slice(available, buf.len());
            let dst = self.buf.mut_slice_from(0);
            std::vec::bytes::copy_memory(dst, src, size);
            self.pos = size;
        } else {
            //# Writing so much data that we might as well write
            //# directly to avoid a copy.
            self.inner.write(self.buf.slice(0, self.pos));
            self.pos = 0;
            self.inner.write(buf);
        }
    }

    fn flush(&mut self) {
        if (self.pos > 0) {
            self.inner.write(self.buf.slice(0, self.pos));
            self.pos = 0;
        }
    }
}
<|endoftext|>


<commit_before>pub mod doc_id_set;
pub mod statistics;
pub mod planner;

use kite::query::Query;
use kite::collectors::{Collector, DocumentMatch};
use byteorder::{ByteOrder, BigEndian};

use key_builder::KeyBuilder;
use document_index::DocRef;
use super::RocksDBIndexReader;
use search::doc_id_set::DocIdSet;
use search::statistics::StatisticsReader;
use search::planner::{SearchPlan, plan_query};
use search::planner::boolean_query::BooleanQueryOp;
use search::planner::score_function::{CombinatorScorer, ScoreFunctionOp};

impl<'a> RocksDBIndexReader<'a> {
    fn search_chunk_boolean_phase(&self, plan: &SearchPlan, chunk: u32) -> DocIdSet {
        // Execute boolean query
        let mut stack = Vec::new();
        for op in plan.boolean_query.iter() {
            match *op {
                BooleanQueryOp::PushEmpty => {
                    stack.push(DocIdSet::new_filled(0));
                }
                BooleanQueryOp::PushFull => {
                    stack.push(DocIdSet::new_filled(65536));
                }
                BooleanQueryOp::PushTermDirectory(field_ref, term_ref) => {
                    let kb = KeyBuilder::chunk_dir_list(chunk, field_ref.ord(), term_ref.ord());
                    match self.snapshot.get(&kb.key()) {
                        Ok(Some(docid_set)) => {
                            let data = DocIdSet::FromRDB(docid_set);
                            stack.push(data);
                        }
                        Ok(None) => stack.push(DocIdSet::new_filled(0)),
                        Err(e) => {},  // FIXME
                    }
                }
                BooleanQueryOp::PushDeletionList => {
                    let kb = KeyBuilder::chunk_del_list(chunk);
                    match self.snapshot.get(&kb.key()) {
                        Ok(Some(deletion_list)) => {
                            let data = DocIdSet::FromRDB(deletion_list);
                            stack.push(data);
                        }
                        Ok(None) => stack.push(DocIdSet::new_filled(0)),
                        Err(e) => {},  // FIXME
                    }
                }
                BooleanQueryOp::And => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.intersection(&b));
                }
                BooleanQueryOp::Or => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.union(&b));
                }
                BooleanQueryOp::AndNot => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.exclusion(&b));
                }
            }
        }

        if !stack.len() == 1 {
            // TODO: Error
        }
        let mut matches = stack.pop().unwrap();

        // Invert the list if the query is negated
        if plan.boolean_query_is_negated {
            let kb = KeyBuilder::chunk_stat(chunk, b"total_docs");
            let total_docs = match self.snapshot.get(&kb.key()) {
                Ok(Some(total_docs)) => BigEndian::read_i64(&total_docs) as u32,
                Ok(None) => 0,
                Err(e) => 0,  // FIXME
            };

            let all_docs = DocIdSet::new_filled(total_docs);
            matches = all_docs.exclusion(&matches);
        }

        matches
    }

    fn score_doc(&self, doc_id: u16, plan: &SearchPlan, chunk: u32, mut stats: &mut StatisticsReader) -> f64 {
        // Execute score function
        let mut stack = Vec::new();
        for op in plan.score_function.iter() {
            match *op {
                ScoreFunctionOp::Literal(val) => stack.push(val),
                ScoreFunctionOp::TermScorer(field_ref, term_ref, ref scorer) => {
                    let kb = KeyBuilder::chunk_dir_list(chunk, field_ref.ord(), term_ref.ord());
                    match self.snapshot.get(&kb.key()) {
                        Ok(Some(data)) => {
                            let docid_set = DocIdSet::FromRDB(data);

                            if docid_set.contains_doc(doc_id) {
                                // Read field length
                                // TODO: we only need this for BM25
                                let kb = KeyBuilder::stored_field_value(chunk, doc_id, field_ref.ord(), b"len");
                                let field_length = match self.snapshot.get(&kb.key()) {
                                    Ok(Some(value)) => {
                                        let length_sqrt = (value[0] as f64) / 3.0 + 1.0;
                                        length_sqrt * length_sqrt
                                    }
                                    Ok(None) => 1.0,
                                    Err(e) => 1.0,  // TODO Error
                                };

                                // Read term frequency
                                let mut value_type = vec![b't', b'f'];
                                value_type.extend(term_ref.ord().to_string().as_bytes());
                                let kb = KeyBuilder::stored_field_value(chunk, doc_id, field_ref.ord(), &value_type);
                                let term_frequency = match self.snapshot.get(&kb.key()) {
                                    Ok(Some(value)) => BigEndian::read_i64(&value),
                                    Ok(None) => 1,
                                    Err(e) => 1,  // TODO Error
                                };

                                let score = scorer.similarity_model.score(1, field_length, stats.total_tokens(field_ref) as u64, stats.total_docs(field_ref) as u64, stats.term_document_frequency(field_ref, term_ref) as u64);
                                stack.push(score * scorer.boost);
                            } else {
                                stack.push(0.0f64);
                            }
                        }
                        Ok(None) => stack.push(0.0f64),
                        Err(e) => {},  // FIXME
                    }
                }
                ScoreFunctionOp::CombinatorScorer(num_vals, ref scorer) => {
                    let score = match *scorer {
                        CombinatorScorer::Avg => {
                            let mut total_score = 0.0f64;

                            for _ in 0..num_vals {
                                total_score += stack.pop().expect("stack underflow");
                            }

                            total_score / num_vals as f64
                        }
                        CombinatorScorer::Max => {
                            let mut max_score = 0.0f64;

                            for _ in 0..num_vals {
                                let score = stack.pop().expect("stack underflow");
                                if score > max_score {
                                    max_score = score
                                }
                            }

                            max_score
                        }
                    };

                    stack.push(score);
                }
            }
        }

        stack.pop().expect("stack underflow")
    }

    fn search_chunk<C: Collector>(&self, collector: &mut C, plan: &SearchPlan, chunk: u32, mut stats: &mut StatisticsReader) {
        let matches = self.search_chunk_boolean_phase(plan, chunk);

        // Score documents and pass to collector
        for doc in matches.iter() {
            let score = self.score_doc(doc, plan, chunk, &mut stats);

            let doc_ref = DocRef::from_chunk_ord(chunk, doc);
            let doc_match = DocumentMatch::new_scored(doc_ref.as_u64(), score);
            collector.collect(doc_match);
        }
    }

    pub fn search<C: Collector>(&self, collector: &mut C, query: &Query) {
        // Plan query
        let plan = plan_query(&self, query, collector.needs_score());

        // Initialise statistics reader
        let mut stats = StatisticsReader::new(&self);

        // Run query on each chunk
        for chunk in self.store.chunks.iter_active(&self.snapshot) {
            self.search_chunk(collector, &plan, chunk, &mut stats);
        }
    }
}
<commit_msg>Only pass in what's needed to query executors<commit_after>pub mod doc_id_set;
pub mod statistics;
pub mod planner;

use kite::query::Query;
use kite::collectors::{Collector, DocumentMatch};
use byteorder::{ByteOrder, BigEndian};

use key_builder::KeyBuilder;
use document_index::DocRef;
use super::RocksDBIndexReader;
use search::doc_id_set::DocIdSet;
use search::statistics::StatisticsReader;
use search::planner::{SearchPlan, plan_query};
use search::planner::boolean_query::BooleanQueryOp;
use search::planner::score_function::{CombinatorScorer, ScoreFunctionOp};

impl<'a> RocksDBIndexReader<'a> {
    fn search_chunk_boolean_phase(&self, boolean_query: &Vec<BooleanQueryOp>, is_negated: bool, chunk: u32) -> DocIdSet {
        // Execute boolean query
        let mut stack = Vec::new();
        for op in boolean_query.iter() {
            match *op {
                BooleanQueryOp::PushEmpty => {
                    stack.push(DocIdSet::new_filled(0));
                }
                BooleanQueryOp::PushFull => {
                    stack.push(DocIdSet::new_filled(65536));
                }
                BooleanQueryOp::PushTermDirectory(field_ref, term_ref) => {
                    let kb = KeyBuilder::chunk_dir_list(chunk, field_ref.ord(), term_ref.ord());
                    match self.snapshot.get(&kb.key()) {
                        Ok(Some(docid_set)) => {
                            let data = DocIdSet::FromRDB(docid_set);
                            stack.push(data);
                        }
                        Ok(None) => stack.push(DocIdSet::new_filled(0)),
                        Err(e) => {},  // FIXME
                    }
                }
                BooleanQueryOp::PushDeletionList => {
                    let kb = KeyBuilder::chunk_del_list(chunk);
                    match self.snapshot.get(&kb.key()) {
                        Ok(Some(deletion_list)) => {
                            let data = DocIdSet::FromRDB(deletion_list);
                            stack.push(data);
                        }
                        Ok(None) => stack.push(DocIdSet::new_filled(0)),
                        Err(e) => {},  // FIXME
                    }
                }
                BooleanQueryOp::And => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.intersection(&b));
                }
                BooleanQueryOp::Or => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.union(&b));
                }
                BooleanQueryOp::AndNot => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.exclusion(&b));
                }
            }
        }

        if !stack.len() == 1 {
            // TODO: Error
        }
        let mut matches = stack.pop().unwrap();

        // Invert the list if the query is negated
        if is_negated {
            let kb = KeyBuilder::chunk_stat(chunk, b"total_docs");
            let total_docs = match self.snapshot.get(&kb.key()) {
                Ok(Some(total_docs)) => BigEndian::read_i64(&total_docs) as u32,
                Ok(None) => 0,
                Err(e) => 0,  // FIXME
            };

            let all_docs = DocIdSet::new_filled(total_docs);
            matches = all_docs.exclusion(&matches);
        }

        matches
    }

    fn score_doc(&self, doc_id: u16, score_function: &Vec<ScoreFunctionOp>, chunk: u32, mut stats: &mut StatisticsReader) -> f64 {
        // Execute score function
        let mut stack = Vec::new();
        for op in score_function.iter() {
            match *op {
                ScoreFunctionOp::Literal(val) => stack.push(val),
                ScoreFunctionOp::TermScorer(field_ref, term_ref, ref scorer) => {
                    let kb = KeyBuilder::chunk_dir_list(chunk, field_ref.ord(), term_ref.ord());
                    match self.snapshot.get(&kb.key()) {
                        Ok(Some(data)) => {
                            let docid_set = DocIdSet::FromRDB(data);

                            if docid_set.contains_doc(doc_id) {
                                // Read field length
                                // TODO: we only need this for BM25
                                let kb = KeyBuilder::stored_field_value(chunk, doc_id, field_ref.ord(), b"len");
                                let field_length = match self.snapshot.get(&kb.key()) {
                                    Ok(Some(value)) => {
                                        let length_sqrt = (value[0] as f64) / 3.0 + 1.0;
                                        length_sqrt * length_sqrt
                                    }
                                    Ok(None) => 1.0,
                                    Err(e) => 1.0,  // TODO Error
                                };

                                // Read term frequency
                                let mut value_type = vec![b't', b'f'];
                                value_type.extend(term_ref.ord().to_string().as_bytes());
                                let kb = KeyBuilder::stored_field_value(chunk, doc_id, field_ref.ord(), &value_type);
                                let term_frequency = match self.snapshot.get(&kb.key()) {
                                    Ok(Some(value)) => BigEndian::read_i64(&value),
                                    Ok(None) => 1,
                                    Err(e) => 1,  // TODO Error
                                };

                                let score = scorer.similarity_model.score(1, field_length, stats.total_tokens(field_ref) as u64, stats.total_docs(field_ref) as u64, stats.term_document_frequency(field_ref, term_ref) as u64);
                                stack.push(score * scorer.boost);
                            } else {
                                stack.push(0.0f64);
                            }
                        }
                        Ok(None) => stack.push(0.0f64),
                        Err(e) => {},  // FIXME
                    }
                }
                ScoreFunctionOp::CombinatorScorer(num_vals, ref scorer) => {
                    let score = match *scorer {
                        CombinatorScorer::Avg => {
                            let mut total_score = 0.0f64;

                            for _ in 0..num_vals {
                                total_score += stack.pop().expect("stack underflow");
                            }

                            total_score / num_vals as f64
                        }
                        CombinatorScorer::Max => {
                            let mut max_score = 0.0f64;

                            for _ in 0..num_vals {
                                let score = stack.pop().expect("stack underflow");
                                if score > max_score {
                                    max_score = score
                                }
                            }

                            max_score
                        }
                    };

                    stack.push(score);
                }
            }
        }

        stack.pop().expect("stack underflow")
    }

    fn search_chunk<C: Collector>(&self, collector: &mut C, plan: &SearchPlan, chunk: u32, mut stats: &mut StatisticsReader) {
        let matches = self.search_chunk_boolean_phase(&plan.boolean_query, plan.boolean_query_is_negated, chunk);

        // Score documents and pass to collector
        for doc in matches.iter() {
            let score = self.score_doc(doc, &plan.score_function, chunk, &mut stats);

            let doc_ref = DocRef::from_chunk_ord(chunk, doc);
            let doc_match = DocumentMatch::new_scored(doc_ref.as_u64(), score);
            collector.collect(doc_match);
        }
    }

    pub fn search<C: Collector>(&self, collector: &mut C, query: &Query) {
        // Plan query
        let plan = plan_query(&self, query, collector.needs_score());

        // Initialise statistics reader
        let mut stats = StatisticsReader::new(&self);

        // Run query on each chunk
        for chunk in self.store.chunks.iter_active(&self.snapshot) {
            self.search_chunk(collector, &plan, chunk, &mut stats);
        }
    }
}
<|endoftext|>