<commit_before>#![allow(dead_code)]
use ffmpeg_sys;
use libc::*;
use std::cmp;
use std::ffi::{ CStr, CString };
use std::ptr;

use errors::*;

#[derive(Debug, Clone, Copy)]
pub struct Rational {
    pub numerator: i32,
    pub denominator: i32,
}

#[derive(Clone, Copy)]
pub struct Codec {
    ptr: *mut ffmpeg_sys::AVCodec,
}

unsafe impl Send for Codec {}
unsafe impl Sync for Codec {}

pub struct PixelFormats {
    codec: Codec,
    index: isize,
}

pub struct Context {
    // The Codec which was used to initialize this Context.
    codec: Codec,
    ptr: *mut ffmpeg_sys::AVCodecContext,
}

pub struct OpenContext {
    context: Context,
}

pub struct Frame {
    ptr: *mut ffmpeg_sys::AVFrame,
}

pub struct OutputContext {
    ptr: *mut ffmpeg_sys::AVFormatContext,
}

impl From<(i32, i32)> for Rational {
    fn from((numerator, denominator): (i32, i32)) -> Self {
        Self {
            numerator,
            denominator,
        }
    }
}

impl From<ffmpeg_sys::AVRational> for Rational {
    fn from(rational: ffmpeg_sys::AVRational) -> Self {
        Self {
            numerator: rational.num,
            denominator: rational.den,
        }
    }
}

impl From<Rational> for ffmpeg_sys::AVRational {
    fn from(rational: Rational) -> Self {
        Self {
            num: rational.numerator,
            den: rational.denominator,
        }
    }
}

impl Codec {
    pub fn description(self) -> String {
        unsafe {
            CStr::from_ptr((*self.ptr).long_name).to_string_lossy().into_owned()
        }
    }

    #[inline]
    pub fn kind(self) -> ffmpeg_sys::AVMediaType {
        unsafe {
            (*self.ptr).kind
        }
    }

    #[inline]
    pub fn is_video(self) -> bool {
        self.kind() == ffmpeg_sys::AVMediaType::AVMEDIA_TYPE_VIDEO
    }

    pub fn pixel_formats(self) -> Option<PixelFormats> {
        unsafe {
            if (*self.ptr).pix_fmts.is_null() {
                None
            } else {
                Some(PixelFormats::new(self))
            }
        }
    }

    pub fn context(self) -> Result<Context> {
        let ptr = unsafe {
            ffmpeg_sys::avcodec_alloc_context3(self.ptr)
        };

        ensure!(!ptr.is_null(), "unable to allocate the codec context");

        Ok(Context {
            codec: self,
            ptr,
        })
    }
}

impl PixelFormats {
    fn new(codec: Codec) -> Self {
        Self {
            codec,
            index: 0,
        }
    }
}

impl Iterator for PixelFormats {
    type Item = ffmpeg_sys::AVPixelFormat;

    fn next(&mut self) -> Option<Self::Item> {
        let format = unsafe {
            *(*self.codec.ptr).pix_fmts.offset(self.index)
        };

        if format == ffmpeg_sys::AVPixelFormat::AV_PIX_FMT_NONE {
            None
        } else {
            self.index += 1;
            Some(format)
        }
    }
}

impl Context {
    #[inline]
    pub fn width(&self) -> u32 {
        unsafe {
            cmp::max(0, (*self.ptr).width) as u32
        }
    }

    #[inline]
    pub fn set_width(&mut self, width: u32) {
        unsafe {
            (*self.ptr).width = cmp::min(width, c_int::max_value() as u32) as c_int;
        }
    }

    #[inline]
    pub fn height(&self) -> u32 {
        unsafe {
            cmp::max(0, (*self.ptr).height) as u32
        }
    }

    #[inline]
    pub fn set_height(&mut self, height: u32) {
        unsafe {
            (*self.ptr).height = cmp::min(height, c_int::max_value() as u32) as c_int;
        }
    }

    #[inline]
    pub fn time_base(&self) -> Rational {
        unsafe {
            (*self.ptr).time_base.into()
        }
    }

    #[inline]
    pub fn set_time_base(&mut self, time_base: &Rational) {
        unsafe {
            (*self.ptr).time_base = (*time_base).into();
        }
    }

    #[inline]
    pub fn pixel_format(&self) -> ffmpeg_sys::AVPixelFormat {
        unsafe {
            (*self.ptr).pix_fmt
        }
    }

    #[inline]
    pub fn set_pixel_format(&mut self, pixel_format: ffmpeg_sys::AVPixelFormat) {
        unsafe {
            (*self.ptr).pix_fmt = pixel_format;
        }
    }

    pub fn open(self) -> Result<OpenContext> {
        let rv = unsafe {
            ffmpeg_sys::avcodec_open2(self.ptr, self.codec.ptr, ptr::null_mut())
        };

        ensure!(rv == 0, "error opening context");

        Ok(OpenContext { context: self })
    }
}

impl Drop for Context {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::avcodec_free_context(&mut self.ptr);
        }
    }
}

impl Frame {
    pub fn new(pixel_format: ffmpeg_sys::AVPixelFormat,
               width: u32,
               height: u32) -> Result<Self> {
        let ptr = unsafe {
            ffmpeg_sys::av_frame_alloc()
        };

        ensure!(!ptr.is_null(), "unable to allocate a frame");

        unsafe {
            (*ptr).format = pixel_format as c_int;
            (*ptr).width = cmp::min(width, c_int::max_value() as u32) as c_int;
            (*ptr).height = cmp::min(height, c_int::max_value() as u32) as c_int;
        }

        Ok(Self { ptr })
    }
}

impl Drop for Frame {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::av_frame_free(&mut self.ptr);
        }
    }
}

impl OutputContext {
    pub fn new(filename: &str) -> Result<Self> {
        let filename = CString::new(filename)
            .chain_err(|| "unable to convert filename to a CString")?;

        let mut ptr = ptr::null_mut();

        let rv = unsafe {
            ffmpeg_sys::avformat_alloc_output_context2(&mut ptr,
                                                       ptr::null_mut(),
                                                       ptr::null_mut(),
                                                       filename.as_ptr())
        };

        // TODO: check and report the error code.
        ensure!(rv >= 0, "unable to allocate the output context");

        Ok(Self { ptr })
    }
}

impl Drop for OutputContext {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::avformat_free_context(self.ptr);
        }
    }
}

pub fn initialize() {
    unsafe {
        ffmpeg_sys::av_register_all();
        ffmpeg_sys::avcodec_register_all();
    }
}

pub fn find_encoder_by_name(name: &str) -> Result<Option<Codec>> {
    let name = CString::new(name).chain_err(|| "could not convert name to CString")?;

    let codec = unsafe {
        ffmpeg_sys::avcodec_find_encoder_by_name(name.as_ptr())
    };

    if codec.is_null() {
        Ok(None)
    } else {
        Ok(Some(Codec { ptr: codec }))
    }
}
<commit_msg>Mark avcodec fields as pub for the time being<commit_after>#![allow(dead_code)]
use ffmpeg_sys;
use libc::*;
use std::cmp;
use std::ffi::{ CStr, CString };
use std::ptr;

use errors::*;

#[derive(Debug, Clone, Copy)]
pub struct Rational {
    pub numerator: i32,
    pub denominator: i32,
}

#[derive(Clone, Copy)]
pub struct Codec {
    ptr: *mut ffmpeg_sys::AVCodec,
}

unsafe impl Send for Codec {}
unsafe impl Sync for Codec {}

pub struct PixelFormats {
    codec: Codec,
    index: isize,
}

pub struct Context {
    // The Codec which was used to initialize this Context.
    codec: Codec,
    pub ptr: *mut ffmpeg_sys::AVCodecContext,
}

pub struct OpenContext {
    pub context: Context,
}

pub struct Frame {
    pub ptr: *mut ffmpeg_sys::AVFrame,
}

pub struct OutputContext {
    pub ptr: *mut ffmpeg_sys::AVFormatContext,
}

impl From<(i32, i32)> for Rational {
    fn from((numerator, denominator): (i32, i32)) -> Self {
        Self {
            numerator,
            denominator,
        }
    }
}

impl From<ffmpeg_sys::AVRational> for Rational {
    fn from(rational: ffmpeg_sys::AVRational) -> Self {
        Self {
            numerator: rational.num,
            denominator: rational.den,
        }
    }
}

impl From<Rational> for ffmpeg_sys::AVRational {
    fn from(rational: Rational) -> Self {
        Self {
            num: rational.numerator,
            den: rational.denominator,
        }
    }
}

impl Codec {
    pub fn description(self) -> String {
        unsafe {
            CStr::from_ptr((*self.ptr).long_name).to_string_lossy().into_owned()
        }
    }

    #[inline]
    pub fn kind(self) -> ffmpeg_sys::AVMediaType {
        unsafe {
            (*self.ptr).kind
        }
    }

    #[inline]
    pub fn is_video(self) -> bool {
        self.kind() == ffmpeg_sys::AVMediaType::AVMEDIA_TYPE_VIDEO
    }

    pub fn pixel_formats(self) -> Option<PixelFormats> {
        unsafe {
            if (*self.ptr).pix_fmts.is_null() {
                None
            } else {
                Some(PixelFormats::new(self))
            }
        }
    }

    pub fn context(self) -> Result<Context> {
        let ptr = unsafe {
            ffmpeg_sys::avcodec_alloc_context3(self.ptr)
        };

        ensure!(!ptr.is_null(), "unable to allocate the codec context");

        Ok(Context {
            codec: self,
            ptr,
        })
    }
}

impl PixelFormats {
    fn new(codec: Codec) -> Self {
        Self {
            codec,
            index: 0,
        }
    }
}

impl Iterator for PixelFormats {
    type Item = ffmpeg_sys::AVPixelFormat;

    fn next(&mut self) -> Option<Self::Item> {
        let format = unsafe {
            *(*self.codec.ptr).pix_fmts.offset(self.index)
        };

        if format == ffmpeg_sys::AVPixelFormat::AV_PIX_FMT_NONE {
            None
        } else {
            self.index += 1;
            Some(format)
        }
    }
}

impl Context {
    #[inline]
    pub fn width(&self) -> u32 {
        unsafe {
            cmp::max(0, (*self.ptr).width) as u32
        }
    }

    #[inline]
    pub fn set_width(&mut self, width: u32) {
        unsafe {
            (*self.ptr).width = cmp::min(width, c_int::max_value() as u32) as c_int;
        }
    }

    #[inline]
    pub fn height(&self) -> u32 {
        unsafe {
            cmp::max(0, (*self.ptr).height) as u32
        }
    }

    #[inline]
    pub fn set_height(&mut self, height: u32) {
        unsafe {
            (*self.ptr).height = cmp::min(height, c_int::max_value() as u32) as c_int;
        }
    }

    #[inline]
    pub fn time_base(&self) -> Rational {
        unsafe {
            (*self.ptr).time_base.into()
        }
    }

    #[inline]
    pub fn set_time_base(&mut self, time_base: &Rational) {
        unsafe {
            (*self.ptr).time_base = (*time_base).into();
        }
    }

    #[inline]
    pub fn pixel_format(&self) -> ffmpeg_sys::AVPixelFormat {
        unsafe {
            (*self.ptr).pix_fmt
        }
    }

    #[inline]
    pub fn set_pixel_format(&mut self, pixel_format: ffmpeg_sys::AVPixelFormat) {
        unsafe {
            (*self.ptr).pix_fmt = pixel_format;
        }
    }

    pub fn open(self) -> Result<OpenContext> {
        let rv = unsafe {
            ffmpeg_sys::avcodec_open2(self.ptr, self.codec.ptr, ptr::null_mut())
        };

        ensure!(rv == 0, "error opening context");

        Ok(OpenContext { context: self })
    }
}

impl Drop for Context {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::avcodec_free_context(&mut self.ptr);
        }
    }
}

impl Frame {
    pub fn new(pixel_format: ffmpeg_sys::AVPixelFormat,
               width: u32,
               height: u32) -> Result<Self> {
        let ptr = unsafe {
            ffmpeg_sys::av_frame_alloc()
        };

        ensure!(!ptr.is_null(), "unable to allocate a frame");

        unsafe {
            (*ptr).format = pixel_format as c_int;
            (*ptr).width = cmp::min(width, c_int::max_value() as u32) as c_int;
            (*ptr).height = cmp::min(height, c_int::max_value() as u32) as c_int;
        }

        Ok(Self { ptr })
    }
}

impl Drop for Frame {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::av_frame_free(&mut self.ptr);
        }
    }
}

impl OutputContext {
    pub fn new(filename: &str) -> Result<Self> {
        let filename = CString::new(filename)
            .chain_err(|| "unable to convert filename to a CString")?;

        let mut ptr = ptr::null_mut();

        let rv = unsafe {
            ffmpeg_sys::avformat_alloc_output_context2(&mut ptr,
                                                       ptr::null_mut(),
                                                       ptr::null_mut(),
                                                       filename.as_ptr())
        };

        // TODO: check and report the error code.
        ensure!(rv >= 0, "unable to allocate the output context");

        Ok(Self { ptr })
    }
}

impl Drop for OutputContext {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::avformat_free_context(self.ptr);
        }
    }
}

pub fn initialize() {
    unsafe {
        ffmpeg_sys::av_register_all();
        ffmpeg_sys::avcodec_register_all();
    }
}

pub fn find_encoder_by_name(name: &str) -> Result<Option<Codec>> {
    let name = CString::new(name).chain_err(|| "could not convert name to CString")?;

    let codec = unsafe {
        ffmpeg_sys::avcodec_find_encoder_by_name(name.as_ptr())
    };

    if codec.is_null() {
        Ok(None)
    } else {
        Ok(Some(Codec { ptr: codec }))
    }
}
<|endoftext|>


<commit_before>use std::path::PathBuf;
use std::process::Command;
use std::env;
use std::io::stderr;
use std::io::Write;

pub use clap::App;

use clap::{Arg, ArgMatches};
use log;
use log::LogLevelFilter;

use configuration::Configuration;
use error::RuntimeError;
use error::RuntimeErrorKind;
use logger::ImagLogger;

use libimagstore::store::Store;

#[derive(Debug)]
pub struct Runtime<'a> {
    rtp: PathBuf,
    configuration: Option<Configuration>,
    cli_matches: ArgMatches<'a>,
    store: Store,
}

impl<'a> Runtime<'a> {

    /**
     * Gets the CLI spec for the program and retreives the config file path (or uses the default on
     * in $HOME/.imag/config, $XDG_CONFIG_DIR/imag/config or from env("$IMAG_CONFIG")
     * and builds the Runtime object with it.
     *
     * The cli_spec object should be initially build with the ::get_default_cli_builder() function.
     *
     */
    pub fn new(cli_spec: App<'a, 'a>) -> Result<Runtime<'a>, RuntimeError> {
        use std::env;

        use libimagstore::hook::position::HookPosition;
        use libimagstore::error::StoreErrorKind;
        use libimagstorestdhook::debug::DebugHook;
        use libimagerror::trace::trace_error;
        use libimagerror::trace::trace_error_dbg;
        use libimagerror::into::IntoError;

        use configuration::error::ConfigErrorKind;

        let matches = cli_spec.get_matches();

        let is_debugging = matches.is_present("debugging");
        let is_verbose   = matches.is_present("verbosity");

        Runtime::init_logger(is_debugging, is_verbose);

        let rtp : PathBuf = matches.value_of("runtimepath")
            .map_or_else(|| {
                env::var("HOME")
                    .map(PathBuf::from)
                    .map(|mut p| { p.push(".imag"); p})
                    .unwrap_or_else(|_| {
                        panic!("You seem to be $HOME-less. Please get a $HOME before using this software. We are sorry for you and hope you have some accommodation anyways.");
                    })
            }, PathBuf::from);
        let storepath = matches.value_of("storepath")
                                .map_or_else(|| {
                                    let mut spath = rtp.clone();
                                    spath.push("store");
                                    spath
                                }, PathBuf::from);

        let configpath = matches.value_of("config")
                                .map_or_else(|| {
                                    let mut spath = rtp.clone();
                                    spath.push("store");
                                    spath
                                }, PathBuf::from);

        let cfg = match Configuration::new(&configpath) {
            Err(e) => if e.err_type() != ConfigErrorKind::NoConfigFileFound {
                return Err(RuntimeErrorKind::Instantiate.into_error_with_cause(Box::new(e)));
            } else {
                warn!("No config file found.");
                warn!("Continuing without configuration file");
                None
            },

            Ok(cfg) => Some(cfg),
        };

        let store_config = match cfg {
            Some(ref c) => c.store_config().cloned(),
            None        => None,
        };

        if is_debugging {
            write!(stderr(), "Config: {:?}\n", cfg).ok();
            write!(stderr(), "Store-config: {:?}\n", store_config).ok();
        }

        Store::new(storepath, store_config).map(|mut store| {
            // If we are debugging, generate hooks for all positions
            if is_debugging {
                let hooks = vec![
                    (DebugHook::new(HookPosition::PreCreate), HookPosition::PreCreate),
                    (DebugHook::new(HookPosition::PostCreate), HookPosition::PostCreate),
                    (DebugHook::new(HookPosition::PreRetrieve), HookPosition::PreRetrieve),
                    (DebugHook::new(HookPosition::PostRetrieve), HookPosition::PostRetrieve),
                    (DebugHook::new(HookPosition::PreUpdate), HookPosition::PreUpdate),
                    (DebugHook::new(HookPosition::PostUpdate), HookPosition::PostUpdate),
                    (DebugHook::new(HookPosition::PreDelete), HookPosition::PreDelete),
                    (DebugHook::new(HookPosition::PostDelete), HookPosition::PostDelete),
                ];

                // Put all debug hooks into the aspect "debug".
                // If it fails, trace the error and warn, but continue.
                for (hook, position) in hooks {
                    if let Err(e) = store.register_hook(position, &String::from("debug"), Box::new(hook)) {
                        if e.err_type() == StoreErrorKind::HookRegisterError {
                            trace_error_dbg(&e);
                            warn!("Registering debug hook with store failed");
                        } else {
                            trace_error(&e);
                        };
                    }
                }
            }

            Runtime {
                cli_matches: matches,
                configuration: cfg,
                rtp: rtp,
                store: store,
            }
        })
        .map_err(Box::new)
        .map_err(|e| RuntimeErrorKind::Instantiate.into_error_with_cause(e))
    }

    /**
     * Get a commandline-interface builder object from `clap`
     *
     * This commandline interface builder object already contains some predefined interface flags:
     *   * -v | --verbose for verbosity
     *   * --debug for debugging
     *   * -c <file> | --config <file> for alternative configuration file
     *   * -r <path> | --rtp <path> for alternative runtimepath
     *   * --store <path> for alternative store path
     * Each has the appropriate help text included.
     *
     * The `appname` shall be "imag-<command>".
     */
    pub fn get_default_cli_builder(appname: &'a str,
                                   version: &'a str,
                                   about: &'a str)
        -> App<'a, 'a>
    {
        App::new(appname)
            .version(version)
            .author("Matthias Beyer <mail@beyermatthias.de>")
            .about(about)
            .arg(Arg::with_name("verbosity")
                .short("v")
                .long("verbose")
                .help("Enables verbosity")
                .required(false)
                .takes_value(false))

            .arg(Arg::with_name("debugging")
                .long("debug")
                .help("Enables debugging output")
                .required(false)
                .takes_value(false))

            .arg(Arg::with_name("config")
                .long("config")
                .help("Path to alternative config file")
                .required(false)
                .takes_value(true))

            .arg(Arg::with_name("runtimepath")
                .long("rtp")
                .help("Alternative runtimepath")
                .required(false)
                .takes_value(true))

            .arg(Arg::with_name("storepath")
                .long("store")
                .help("Alternative storepath. Must be specified as full path, can be outside of the RTP")
                .required(false)
                .takes_value(true))

            .arg(Arg::with_name("editor")
                .long("editor")
                .help("Set editor")
                .required(false)
                .takes_value(true))
    }

    /**
     * Initialize the internal logger
     */
    fn init_logger(is_debugging: bool, is_verbose: bool) {
        use std::env::var as env_var;
        use env_logger;

        if env_var("IMAG_LOG_ENV").is_ok() {
            env_logger::init().unwrap();
        } else {
            let lvl = if is_debugging {
                LogLevelFilter::Debug
            } else if is_verbose {
                LogLevelFilter::Info
            } else {
                LogLevelFilter::Error
            };

            log::set_logger(|max_log_lvl| {
                max_log_lvl.set(lvl);
                debug!("Init logger with {}", lvl);
                Box::new(ImagLogger::new(lvl.to_log_level().unwrap()))
            })
            .map_err(|_| {
                panic!("Could not setup logger");
            })
            .ok();
        }
    }

    /**
     * Get the verbosity flag value
     */
    pub fn is_verbose(&self) -> bool {
        self.cli_matches.is_present("verbosity")
    }

    /**
     * Get the debugging flag value
     */
    pub fn is_debugging(&self) -> bool {
        self.cli_matches.is_present("debugging")
    }

    /**
     * Get the runtimepath
     */
    pub fn rtp(&self) -> &PathBuf {
        &self.rtp
    }

    /**
     * Get the commandline interface matches
     */
    pub fn cli(&self) -> &ArgMatches {
        &self.cli_matches
    }

    /**
     * Get the configuration object
     */
    pub fn config(&self) -> Option<&Configuration> {
        self.configuration.as_ref()
    }

    /**
     * Get the store object
     */
    pub fn store(&self) -> &Store {
        &self.store
    }

    pub fn editor(&self) -> Option<Command> {
        self.cli()
            .value_of("editor")
            .map(String::from)
            .or({
                match self.configuration {
                    Some(ref c) => c.editor().cloned(),
                    _ => None,
                }
            })
            .or(env::var("EDITOR").ok())
            .map(Command::new)
    }
}

<commit_msg>Replace error boxing and wrapping by helper generated with libimagerror<commit_after>use std::path::PathBuf;
use std::process::Command;
use std::env;
use std::io::stderr;
use std::io::Write;

pub use clap::App;

use clap::{Arg, ArgMatches};
use log;
use log::LogLevelFilter;

use configuration::Configuration;
use error::RuntimeError;
use error::RuntimeErrorKind;
use error::MapErrInto;
use logger::ImagLogger;

use libimagstore::store::Store;

#[derive(Debug)]
pub struct Runtime<'a> {
    rtp: PathBuf,
    configuration: Option<Configuration>,
    cli_matches: ArgMatches<'a>,
    store: Store,
}

impl<'a> Runtime<'a> {

    /**
     * Gets the CLI spec for the program and retreives the config file path (or uses the default on
     * in $HOME/.imag/config, $XDG_CONFIG_DIR/imag/config or from env("$IMAG_CONFIG")
     * and builds the Runtime object with it.
     *
     * The cli_spec object should be initially build with the ::get_default_cli_builder() function.
     *
     */
    pub fn new(cli_spec: App<'a, 'a>) -> Result<Runtime<'a>, RuntimeError> {
        use std::env;

        use libimagstore::hook::position::HookPosition;
        use libimagstore::error::StoreErrorKind;
        use libimagstorestdhook::debug::DebugHook;
        use libimagerror::trace::trace_error;
        use libimagerror::trace::trace_error_dbg;
        use libimagerror::into::IntoError;

        use configuration::error::ConfigErrorKind;

        let matches = cli_spec.get_matches();

        let is_debugging = matches.is_present("debugging");
        let is_verbose   = matches.is_present("verbosity");

        Runtime::init_logger(is_debugging, is_verbose);

        let rtp : PathBuf = matches.value_of("runtimepath")
            .map_or_else(|| {
                env::var("HOME")
                    .map(PathBuf::from)
                    .map(|mut p| { p.push(".imag"); p})
                    .unwrap_or_else(|_| {
                        panic!("You seem to be $HOME-less. Please get a $HOME before using this software. We are sorry for you and hope you have some accommodation anyways.");
                    })
            }, PathBuf::from);
        let storepath = matches.value_of("storepath")
                                .map_or_else(|| {
                                    let mut spath = rtp.clone();
                                    spath.push("store");
                                    spath
                                }, PathBuf::from);

        let configpath = matches.value_of("config")
                                .map_or_else(|| {
                                    let mut spath = rtp.clone();
                                    spath.push("store");
                                    spath
                                }, PathBuf::from);

        let cfg = match Configuration::new(&configpath) {
            Err(e) => if e.err_type() != ConfigErrorKind::NoConfigFileFound {
                return Err(RuntimeErrorKind::Instantiate.into_error_with_cause(Box::new(e)));
            } else {
                warn!("No config file found.");
                warn!("Continuing without configuration file");
                None
            },

            Ok(cfg) => Some(cfg),
        };

        let store_config = match cfg {
            Some(ref c) => c.store_config().cloned(),
            None        => None,
        };

        if is_debugging {
            write!(stderr(), "Config: {:?}\n", cfg).ok();
            write!(stderr(), "Store-config: {:?}\n", store_config).ok();
        }

        Store::new(storepath, store_config).map(|mut store| {
            // If we are debugging, generate hooks for all positions
            if is_debugging {
                let hooks = vec![
                    (DebugHook::new(HookPosition::PreCreate), HookPosition::PreCreate),
                    (DebugHook::new(HookPosition::PostCreate), HookPosition::PostCreate),
                    (DebugHook::new(HookPosition::PreRetrieve), HookPosition::PreRetrieve),
                    (DebugHook::new(HookPosition::PostRetrieve), HookPosition::PostRetrieve),
                    (DebugHook::new(HookPosition::PreUpdate), HookPosition::PreUpdate),
                    (DebugHook::new(HookPosition::PostUpdate), HookPosition::PostUpdate),
                    (DebugHook::new(HookPosition::PreDelete), HookPosition::PreDelete),
                    (DebugHook::new(HookPosition::PostDelete), HookPosition::PostDelete),
                ];

                // Put all debug hooks into the aspect "debug".
                // If it fails, trace the error and warn, but continue.
                for (hook, position) in hooks {
                    if let Err(e) = store.register_hook(position, &String::from("debug"), Box::new(hook)) {
                        if e.err_type() == StoreErrorKind::HookRegisterError {
                            trace_error_dbg(&e);
                            warn!("Registering debug hook with store failed");
                        } else {
                            trace_error(&e);
                        };
                    }
                }
            }

            Runtime {
                cli_matches: matches,
                configuration: cfg,
                rtp: rtp,
                store: store,
            }
        })
        .map_err_into(RuntimeErrorKind::Instantiate)
    }

    /**
     * Get a commandline-interface builder object from `clap`
     *
     * This commandline interface builder object already contains some predefined interface flags:
     *   * -v | --verbose for verbosity
     *   * --debug for debugging
     *   * -c <file> | --config <file> for alternative configuration file
     *   * -r <path> | --rtp <path> for alternative runtimepath
     *   * --store <path> for alternative store path
     * Each has the appropriate help text included.
     *
     * The `appname` shall be "imag-<command>".
     */
    pub fn get_default_cli_builder(appname: &'a str,
                                   version: &'a str,
                                   about: &'a str)
        -> App<'a, 'a>
    {
        App::new(appname)
            .version(version)
            .author("Matthias Beyer <mail@beyermatthias.de>")
            .about(about)
            .arg(Arg::with_name("verbosity")
                .short("v")
                .long("verbose")
                .help("Enables verbosity")
                .required(false)
                .takes_value(false))

            .arg(Arg::with_name("debugging")
                .long("debug")
                .help("Enables debugging output")
                .required(false)
                .takes_value(false))

            .arg(Arg::with_name("config")
                .long("config")
                .help("Path to alternative config file")
                .required(false)
                .takes_value(true))

            .arg(Arg::with_name("runtimepath")
                .long("rtp")
                .help("Alternative runtimepath")
                .required(false)
                .takes_value(true))

            .arg(Arg::with_name("storepath")
                .long("store")
                .help("Alternative storepath. Must be specified as full path, can be outside of the RTP")
                .required(false)
                .takes_value(true))

            .arg(Arg::with_name("editor")
                .long("editor")
                .help("Set editor")
                .required(false)
                .takes_value(true))
    }

    /**
     * Initialize the internal logger
     */
    fn init_logger(is_debugging: bool, is_verbose: bool) {
        use std::env::var as env_var;
        use env_logger;

        if env_var("IMAG_LOG_ENV").is_ok() {
            env_logger::init().unwrap();
        } else {
            let lvl = if is_debugging {
                LogLevelFilter::Debug
            } else if is_verbose {
                LogLevelFilter::Info
            } else {
                LogLevelFilter::Error
            };

            log::set_logger(|max_log_lvl| {
                max_log_lvl.set(lvl);
                debug!("Init logger with {}", lvl);
                Box::new(ImagLogger::new(lvl.to_log_level().unwrap()))
            })
            .map_err(|_| {
                panic!("Could not setup logger");
            })
            .ok();
        }
    }

    /**
     * Get the verbosity flag value
     */
    pub fn is_verbose(&self) -> bool {
        self.cli_matches.is_present("verbosity")
    }

    /**
     * Get the debugging flag value
     */
    pub fn is_debugging(&self) -> bool {
        self.cli_matches.is_present("debugging")
    }

    /**
     * Get the runtimepath
     */
    pub fn rtp(&self) -> &PathBuf {
        &self.rtp
    }

    /**
     * Get the commandline interface matches
     */
    pub fn cli(&self) -> &ArgMatches {
        &self.cli_matches
    }

    /**
     * Get the configuration object
     */
    pub fn config(&self) -> Option<&Configuration> {
        self.configuration.as_ref()
    }

    /**
     * Get the store object
     */
    pub fn store(&self) -> &Store {
        &self.store
    }

    pub fn editor(&self) -> Option<Command> {
        self.cli()
            .value_of("editor")
            .map(String::from)
            .or({
                match self.configuration {
                    Some(ref c) => c.editor().cloned(),
                    _ => None,
                }
            })
            .or(env::var("EDITOR").ok())
            .map(Command::new)
    }
}

<|endoftext|>


<commit_before>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use std::collections::BTreeMap;
use std::error::Error;

use ruru::{Class, Object, AnyObject, Boolean, RString, VM, Hash, NilClass, VerifiedObject};

use libimagstore::store::EntryHeader;
use libimagstore::store::EntryContent;
use libimagstore::store::Entry;
use libimagstore::storeid::StoreId;

use ruby_utils::IntoToml;
use toml_utils::IntoRuby;
use util::Wrap;
use util::Unwrap;
use cache::StoreHandle;

pub struct FileLockEntryHandle(StoreHandle, StoreId);

impl FileLockEntryHandle {
    pub fn new(sh: StoreHandle, id: StoreId) -> FileLockEntryHandle {
        FileLockEntryHandle(sh, id)
    }

    pub fn store_handle(&self) -> &StoreHandle {
        &self.0
    }

    pub fn fle_handle(&self) -> &StoreId {
        &self.1
    }
}

wrappable_struct!(FileLockEntryHandle, FileLockEntryWrapper, FLE_WRAPPER);
class!(RFileLockEntryHandle);
impl_wrap!(FileLockEntryHandle => FLE_WRAPPER);
impl_unwrap!(RFileLockEntryHandle => FileLockEntryHandle => FLE_WRAPPER);
impl_verified_object!(RFileLockEntryHandle);

/// Helper macro for operating on RUBY_STORE_CACHE object
///
/// This helps us calling operations on FileLockEntry objects.
///
/// What I do here: Fetch the Store object from the cache, fetch the appropriate FileLockEntry and
/// call the operation on it.
///
/// This could be improved with another cache, so not the store is cached but the FileLockEntry
/// only, but then we run into lifetime problems with the Store and its FileLockEntry objects.
/// Feel free to fix this, but for now, this is a workable solution.
///
#[macro_export]
macro_rules! call_on_fle_from_store {
    ($itself:ident ($wrapper:ident) -> $name:ident -> $operation:block) => {{
        let handle = $itself.get_data(&*$wrapper);
        let store_handle = handle.store_handle();
        call_on_store_by_handle! {
            store_handle named store inside {
                match store.get(handle.fle_handle().clone()) {
                    Ok(Some(mut $name)) => {
                        $operation
                    },
                    Ok(None) => {
                        VM::raise(Class::from_existing("RuntimeError"), "Obj does not exist");
                        NilClass::new().to_any_object()
                    },
                    Err(e) => {
                        VM::raise(Class::from_existing("RuntimeError"), e.description());
                        NilClass::new().to_any_object()
                    },
                }
            }
        }
    }};
    ($itself:ident ($wrapper:ident) -> $name:ident -> $operation: block on fail return $ex:expr) => {{
        let handle = $itself.get_data(&*$wrapper);
        let store_handle = handle.store_handle();
        call_on_store_by_handle! {
            store_handle named store inside {
                match store.get(handle.fle_handle().clone()) {
                    Ok(Some(mut $name)) => {
                        $operation
                    },
                    Ok(None) => {
                        VM::raise(Class::from_existing("RuntimeError"), "Obj does not exist");
                        $ex
                    },
                    Err(e) => {
                        VM::raise(Class::from_existing("RuntimeError"), e.description());
                        $ex
                    },
                }
            }
        }
    }};
}

methods!(
    RFileLockEntryHandle,
    itself,

    fn r_get_location() -> AnyObject {
        call_on_fle_from_store!(itself (FLE_WRAPPER) -> fle -> { fle.get_location().clone().wrap() })
    }

    fn r_get_header() -> AnyObject {
        call_on_fle_from_store!(itself (FLE_WRAPPER) -> fle -> { fle.get_header().clone().wrap() })
    }

    fn r_set_header(hdr: Hash) -> NilClass {
        use ruby_utils::IntoToml;
        use toml::Value;

        let entryheader = match typecheck!(hdr or return NilClass::new()).into_toml() {
            Value::Table(t) => EntryHeader::from(t),
            _ => {
                let ec = Class::from_existing("RuntimeError");
                VM::raise(ec, "Something weird happened. Hash seems to be not a Hash");
                return NilClass::new();
            },
        };

        call_on_fle_from_store!(itself (FLE_WRAPPER) -> fle -> {
            *fle.get_header_mut() = entryheader;
            NilClass::new().to_any_object()
        });

        NilClass::new()
    }

    fn r_get_content() -> AnyObject {
        call_on_fle_from_store!(itself (FLE_WRAPPER) -> fle -> {
            fle.get_content().clone().wrap()
        } on fail return NilClass::new().to_any_object())
    }

    fn r_set_content(ctt: RString) -> NilClass {
        use ruby_utils::IntoToml;
        use toml::Value;

        let content = match typecheck!(ctt).into_toml() {
            Value::String(s) => s,
            _ => {
                let ec = Class::from_existing("RuntimeError");
                VM::raise(ec, "Something weird happened. String seems to be not a String");
                return NilClass::new();
            },
        };

        call_on_fle_from_store!(itself (FLE_WRAPPER) -> fle -> {
            *fle.get_content_mut() = content;
            NilClass::new().to_any_object()
        });

        NilClass::new()
    }

);

wrappable_struct!(EntryHeader, EntryHeaderWrapper, ENTRY_HEADER_WRAPPER);
class!(REntryHeader);
impl_wrap!(EntryHeader => ENTRY_HEADER_WRAPPER);
impl_unwrap!(REntryHeader => EntryHeader => ENTRY_HEADER_WRAPPER);
impl_verified_object!(REntryHeader);

methods!(
    REntryHeader,
    itself,

    fn r_entry_header_new() -> AnyObject {
        EntryHeader::new().wrap()
    }

    fn r_entry_header_insert(spec: RString, obj: AnyObject) -> Boolean {
        let spec = typecheck!(spec or return Boolean::new(false)).to_string();
        let obj = obj.unwrap(); // possibly not safe... TODO

        match itself.get_data(&*ENTRY_HEADER_WRAPPER).insert(&spec, obj.into_toml()) {
            Ok(b) => Boolean::new(b),
            Err(e) => {
                VM::raise(Class::from_existing("RuntimeError"), e.description());
                Boolean::new(false)
            }
        }
    }

    fn r_entry_header_set(spec: RString, obj: AnyObject) -> AnyObject {
        use ruru::NilClass;

        let spec = typecheck!(spec or return any Boolean::new(false)).to_string();
        let obj = obj.unwrap(); // possibly not safe... TODO

        match itself.get_data(&*ENTRY_HEADER_WRAPPER).set(&spec, obj.into_toml()) {
            Ok(Some(v)) => v.into_ruby(),
            Ok(None)    => NilClass::new().to_any_object(),
            Err(e) => {
                VM::raise(Class::from_existing("RuntimeError"), e.description());
                return Boolean::new(false).to_any_object();
            }
        }
    }

    fn r_entry_header_get(spec: RString) -> AnyObject {
        use ruru::NilClass;

        let spec = typecheck!(spec or return any Boolean::new(false)).to_string();

        match itself.get_data(&*ENTRY_HEADER_WRAPPER).read(&spec) {
            Ok(Some(v)) => v.into_ruby(),
            Ok(None)    => NilClass::new().to_any_object(),
            Err(e) => {
                VM::raise(Class::from_existing("RuntimeError"), e.description());
                return Boolean::new(false).to_any_object();
            }
        }
    }

);

wrappable_struct!(EntryContent, EntryContentWrapper, ENTRY_CONTENT_WRAPPER);
class!(REntryContent);
impl_wrap!(EntryContent => ENTRY_CONTENT_WRAPPER);
impl_unwrap!(REntryContent => EntryContent => ENTRY_CONTENT_WRAPPER);

wrappable_struct!(Entry, EntryWrapper, ENTRY_WRAPPER);
class!(REntry);
impl_unwrap!(REntry => Entry => ENTRY_WRAPPER);

pub fn setup_filelockentry() -> Class {
    let mut class = Class::new("RFileLockEntryHandle", None);
    class.define(|itself| {
        itself.def("location", r_get_location);
        itself.def("header"  , r_get_header);
        itself.def("header=" , r_set_header);
        itself.def("content" , r_get_content);
        itself.def("content=", r_set_content);
    });
    class
}

pub fn setup_entryheader() -> Class {
    let mut class = Class::new("REntryHeader", None);
    class.define(|itself| {
        itself.def("insert", r_entry_header_insert);
        itself.def("set"   , r_entry_header_set);
        itself.def("[]="   , r_entry_header_set);
        itself.def("read"  , r_entry_header_get);
        itself.def("[]"    , r_entry_header_get);
    });
    class
}

pub fn setup_entrycontent() -> Class {
    let string = Class::from_existing("String");
    Class::new("REntryContent", Some(&string))
}
<commit_msg>Fix: Pass error return expression further down<commit_after>//
// imag - the personal information management suite for the commandline
// Copyright (C) 2015, 2016 Matthias Beyer <mail@beyermatthias.de> and contributors
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; version
// 2.1 of the License.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
//

use std::collections::BTreeMap;
use std::error::Error;

use ruru::{Class, Object, AnyObject, Boolean, RString, VM, Hash, NilClass, VerifiedObject};

use libimagstore::store::EntryHeader;
use libimagstore::store::EntryContent;
use libimagstore::store::Entry;
use libimagstore::storeid::StoreId;

use ruby_utils::IntoToml;
use toml_utils::IntoRuby;
use util::Wrap;
use util::Unwrap;
use cache::StoreHandle;

pub struct FileLockEntryHandle(StoreHandle, StoreId);

impl FileLockEntryHandle {
    pub fn new(sh: StoreHandle, id: StoreId) -> FileLockEntryHandle {
        FileLockEntryHandle(sh, id)
    }

    pub fn store_handle(&self) -> &StoreHandle {
        &self.0
    }

    pub fn fle_handle(&self) -> &StoreId {
        &self.1
    }
}

wrappable_struct!(FileLockEntryHandle, FileLockEntryWrapper, FLE_WRAPPER);
class!(RFileLockEntryHandle);
impl_wrap!(FileLockEntryHandle => FLE_WRAPPER);
impl_unwrap!(RFileLockEntryHandle => FileLockEntryHandle => FLE_WRAPPER);
impl_verified_object!(RFileLockEntryHandle);

/// Helper macro for operating on RUBY_STORE_CACHE object
///
/// This helps us calling operations on FileLockEntry objects.
///
/// What I do here: Fetch the Store object from the cache, fetch the appropriate FileLockEntry and
/// call the operation on it.
///
/// This could be improved with another cache, so not the store is cached but the FileLockEntry
/// only, but then we run into lifetime problems with the Store and its FileLockEntry objects.
/// Feel free to fix this, but for now, this is a workable solution.
///
#[macro_export]
macro_rules! call_on_fle_from_store {
    ($itself:ident ($wrapper:ident) -> $name:ident -> $operation:block) => {{
        let handle = $itself.get_data(&*$wrapper);
        let store_handle = handle.store_handle();
        call_on_store_by_handle! {
            store_handle named store inside {
                match store.get(handle.fle_handle().clone()) {
                    Ok(Some(mut $name)) => {
                        $operation
                    },
                    Ok(None) => {
                        VM::raise(Class::from_existing("RuntimeError"), "Obj does not exist");
                        NilClass::new().to_any_object()
                    },
                    Err(e) => {
                        VM::raise(Class::from_existing("RuntimeError"), e.description());
                        NilClass::new().to_any_object()
                    },
                }
            }
        }
    }};
    ($itself:ident ($wrapper:ident) -> $name:ident -> $operation: block on fail return $ex:expr) => {{
        let handle = $itself.get_data(&*$wrapper);
        let store_handle = handle.store_handle();
        call_on_store_by_handle! {
            store_handle named store inside {
                match store.get(handle.fle_handle().clone()) {
                    Ok(Some(mut $name)) => {
                        $operation
                    },
                    Ok(None) => {
                        VM::raise(Class::from_existing("RuntimeError"), "Obj does not exist");
                        $ex
                    },
                    Err(e) => {
                        VM::raise(Class::from_existing("RuntimeError"), e.description());
                        $ex
                    },
                }
            } on fail return $ex
        }
    }};
}

methods!(
    RFileLockEntryHandle,
    itself,

    fn r_get_location() -> AnyObject {
        call_on_fle_from_store!(itself (FLE_WRAPPER) -> fle -> { fle.get_location().clone().wrap() })
    }

    fn r_get_header() -> AnyObject {
        call_on_fle_from_store!(itself (FLE_WRAPPER) -> fle -> { fle.get_header().clone().wrap() })
    }

    fn r_set_header(hdr: Hash) -> NilClass {
        use ruby_utils::IntoToml;
        use toml::Value;

        let entryheader = match typecheck!(hdr or return NilClass::new()).into_toml() {
            Value::Table(t) => EntryHeader::from(t),
            _ => {
                let ec = Class::from_existing("RuntimeError");
                VM::raise(ec, "Something weird happened. Hash seems to be not a Hash");
                return NilClass::new();
            },
        };

        call_on_fle_from_store!(itself (FLE_WRAPPER) -> fle -> {
            *fle.get_header_mut() = entryheader;
            NilClass::new().to_any_object()
        });

        NilClass::new()
    }

    fn r_get_content() -> AnyObject {
        call_on_fle_from_store!(itself (FLE_WRAPPER) -> fle -> {
            fle.get_content().clone().wrap()
        } on fail return NilClass::new().to_any_object())
    }

    fn r_set_content(ctt: RString) -> NilClass {
        use ruby_utils::IntoToml;
        use toml::Value;

        let content = match typecheck!(ctt).into_toml() {
            Value::String(s) => s,
            _ => {
                let ec = Class::from_existing("RuntimeError");
                VM::raise(ec, "Something weird happened. String seems to be not a String");
                return NilClass::new();
            },
        };

        call_on_fle_from_store!(itself (FLE_WRAPPER) -> fle -> {
            *fle.get_content_mut() = content;
            NilClass::new().to_any_object()
        });

        NilClass::new()
    }

);

wrappable_struct!(EntryHeader, EntryHeaderWrapper, ENTRY_HEADER_WRAPPER);
class!(REntryHeader);
impl_wrap!(EntryHeader => ENTRY_HEADER_WRAPPER);
impl_unwrap!(REntryHeader => EntryHeader => ENTRY_HEADER_WRAPPER);
impl_verified_object!(REntryHeader);

methods!(
    REntryHeader,
    itself,

    fn r_entry_header_new() -> AnyObject {
        EntryHeader::new().wrap()
    }

    fn r_entry_header_insert(spec: RString, obj: AnyObject) -> Boolean {
        let spec = typecheck!(spec or return Boolean::new(false)).to_string();
        let obj = obj.unwrap(); // possibly not safe... TODO

        match itself.get_data(&*ENTRY_HEADER_WRAPPER).insert(&spec, obj.into_toml()) {
            Ok(b) => Boolean::new(b),
            Err(e) => {
                VM::raise(Class::from_existing("RuntimeError"), e.description());
                Boolean::new(false)
            }
        }
    }

    fn r_entry_header_set(spec: RString, obj: AnyObject) -> AnyObject {
        use ruru::NilClass;

        let spec = typecheck!(spec or return any Boolean::new(false)).to_string();
        let obj = obj.unwrap(); // possibly not safe... TODO

        match itself.get_data(&*ENTRY_HEADER_WRAPPER).set(&spec, obj.into_toml()) {
            Ok(Some(v)) => v.into_ruby(),
            Ok(None)    => NilClass::new().to_any_object(),
            Err(e) => {
                VM::raise(Class::from_existing("RuntimeError"), e.description());
                return Boolean::new(false).to_any_object();
            }
        }
    }

    fn r_entry_header_get(spec: RString) -> AnyObject {
        use ruru::NilClass;

        let spec = typecheck!(spec or return any Boolean::new(false)).to_string();

        match itself.get_data(&*ENTRY_HEADER_WRAPPER).read(&spec) {
            Ok(Some(v)) => v.into_ruby(),
            Ok(None)    => NilClass::new().to_any_object(),
            Err(e) => {
                VM::raise(Class::from_existing("RuntimeError"), e.description());
                return Boolean::new(false).to_any_object();
            }
        }
    }

);

wrappable_struct!(EntryContent, EntryContentWrapper, ENTRY_CONTENT_WRAPPER);
class!(REntryContent);
impl_wrap!(EntryContent => ENTRY_CONTENT_WRAPPER);
impl_unwrap!(REntryContent => EntryContent => ENTRY_CONTENT_WRAPPER);

wrappable_struct!(Entry, EntryWrapper, ENTRY_WRAPPER);
class!(REntry);
impl_unwrap!(REntry => Entry => ENTRY_WRAPPER);

pub fn setup_filelockentry() -> Class {
    let mut class = Class::new("RFileLockEntryHandle", None);
    class.define(|itself| {
        itself.def("location", r_get_location);
        itself.def("header"  , r_get_header);
        itself.def("header=" , r_set_header);
        itself.def("content" , r_get_content);
        itself.def("content=", r_set_content);
    });
    class
}

pub fn setup_entryheader() -> Class {
    let mut class = Class::new("REntryHeader", None);
    class.define(|itself| {
        itself.def("insert", r_entry_header_insert);
        itself.def("set"   , r_entry_header_set);
        itself.def("[]="   , r_entry_header_set);
        itself.def("read"  , r_entry_header_get);
        itself.def("[]"    , r_entry_header_get);
    });
    class
}

pub fn setup_entrycontent() -> Class {
    let string = Class::from_existing("String");
    Class::new("REntryContent", Some(&string))
}
<|endoftext|>


<commit_before>// # Custom collector example
//
// This example shows how you can implement your own
// collector. As an example, we will compute a collector
// that computes the standard deviation of a given fast field.
//
// Of course, you can have a look at the tantivy's built-in collectors
// such as the `CountCollector` for more examples.

// ---
// Importing tantivy...
use std::marker::PhantomData;

use crate::collector::{Collector, SegmentCollector};
use crate::fastfield::{FastFieldReader, FastValue};
use crate::schema::Field;
use crate::{Score, SegmentReader, TantivyError};

/// The `FilterCollector` collector filters docs using a u64 fast field value and a predicate.
/// Only the documents for which the predicate returned "true" will be passed on to the next collector.
///
/// ```rust
/// use tantivy::collector::{TopDocs, FilterCollector};
/// use tantivy::query::QueryParser;
/// use tantivy::schema::{Schema, FAST, TEXT};
/// use tantivy::DateTime;
/// use std::str::FromStr;
/// use tantivy::{doc, DocAddress, Index};
///
/// let mut schema_builder = Schema::builder();
/// let title = schema_builder.add_text_field("title", TEXT);
/// let price = schema_builder.add_u64_field("price", FAST);
/// let date = schema_builder.add_date_field("date", FAST);
/// let schema = schema_builder.build();
/// let index = Index::create_in_ram(schema);
///
/// let mut index_writer = index.writer_with_num_threads(1, 10_000_000).unwrap();
/// index_writer.add_document(doc!(title => "The Name of the Wind", price => 30_200u64, date => DateTime::from_str("1898-04-09T00:00:00+00:00").unwrap()));
/// index_writer.add_document(doc!(title => "The Diary of Muadib", price => 29_240u64, date => DateTime::from_str("2020-04-09T00:00:00+00:00").unwrap()));
/// index_writer.add_document(doc!(title => "A Dairy Cow", price => 21_240u64, date => DateTime::from_str("2019-04-09T00:00:00+00:00").unwrap()));
/// index_writer.add_document(doc!(title => "The Diary of a Young Girl", price => 20_120u64, date => DateTime::from_str("2018-04-09T00:00:00+00:00").unwrap()));
/// assert!(index_writer.commit().is_ok());
///
/// let reader = index.reader().unwrap();
/// let searcher = reader.searcher();
///
/// let query_parser = QueryParser::for_index(&index, vec![title]);
/// let query = query_parser.parse_query("diary").unwrap();
/// let filter_some_collector = FilterCollector::new(price, &|value: u64| value > 20_120u64, TopDocs::with_limit(2));
/// let top_docs = searcher.search(&query, &filter_some_collector).unwrap();
///
/// assert_eq!(top_docs.len(), 1);
/// assert_eq!(top_docs[0].1, DocAddress(0, 1));
///
/// let filter_all_collector: FilterCollector<_, _, u64> = FilterCollector::new(price, &|value| value < 5u64, TopDocs::with_limit(2));
/// let filtered_top_docs = searcher.search(&query, &filter_all_collector).unwrap();
///
/// assert_eq!(filtered_top_docs.len(), 0);
///
/// fn date_debug(value: DateTime) -> bool {
///     println!("date: {:?}", value);
///     assert_eq!(value, DateTime::from_str("1000-04-09T00:00:00+00:00").unwrap());
///     (value - DateTime::from_str("2019-04-09T00:00:00+00:00").unwrap()).num_weeks() > 0
/// }
///
/// let filter_dates_collector = FilterCollector::new(date, &date_debug, TopDocs::with_limit(2));
/// let filtered_date_docs = searcher.search(&query, &filter_all_collector).unwrap();
///
/// assert_eq!(filtered_date_docs.len(), 5);
/// ```
pub struct FilterCollector<TCollector, TPredicate, TPredicateValue: FastValue>
where
    TPredicate: 'static,
{
    field: Field,
    collector: TCollector,
    predicate: &'static TPredicate,
    t_predicate_value: PhantomData<TPredicateValue>,
}

impl<TCollector, TPredicate, TPredicateValue: FastValue>
    FilterCollector<TCollector, TPredicate, TPredicateValue>
where
    TCollector: Collector + Send + Sync,
    TPredicate: Fn(TPredicateValue) -> bool + Send + Sync,
{
    /// Create a new FilterCollector.
    pub fn new(
        field: Field,
        predicate: &'static TPredicate,
        collector: TCollector,
    ) -> FilterCollector<TCollector, TPredicate, TPredicateValue> {
        FilterCollector {
            field,
            predicate,
            collector,
            t_predicate_value: PhantomData,
        }
    }
}

impl<TCollector, TPredicate, TPredicateValue: FastValue> Collector
    for FilterCollector<TCollector, TPredicate, TPredicateValue>
where
    TCollector: Collector + Send + Sync,
    TPredicate: 'static + Fn(TPredicateValue) -> bool + Send + Sync,
    TPredicateValue: 'static + FastValue,
{
    // That's the type of our result.
    // Our standard deviation will be a float.
    type Fruit = TCollector::Fruit;

    type Child = FilterSegmentCollector<TCollector::Child, TPredicate, TPredicateValue>;

    fn for_segment(
        &self,
        segment_local_id: u32,
        segment_reader: &SegmentReader,
    ) -> crate::Result<FilterSegmentCollector<TCollector::Child, TPredicate, TPredicateValue>> {
        let schema = segment_reader.schema();
        let field_entry = schema.get_field_entry(self.field);
        if !field_entry.is_fast() {
            return Err(TantivyError::SchemaError(format!(
                "Field {:?} is not a fast field.",
                field_entry.name()
            )));
        }
        let requested_type = TPredicateValue::to_type();
        let field_schema_type = field_entry.field_type().value_type();
        if requested_type != field_schema_type {
            return Err(TantivyError::SchemaError(format!(
                "Field {:?} is of type {:?}!={:?}",
                field_entry.name(),
                requested_type,
                field_schema_type
            )));
        }

        let fast_field_reader = segment_reader
            .fast_fields()
            .typed_fast_field_reader(self.field)
            .ok_or_else(|| {
                TantivyError::SchemaError(format!(
                    "{:?} is not declared as a fast field in the schema.",
                    self.field
                ))
            })?;

        let segment_collector = self
            .collector
            .for_segment(segment_local_id, segment_reader)?;

        Ok(FilterSegmentCollector {
            fast_field_reader,
            segment_collector: segment_collector,
            predicate: self.predicate,
            t_predicate_value: PhantomData,
        })
    }

    fn requires_scoring(&self) -> bool {
        self.collector.requires_scoring()
    }

    fn merge_fruits(
        &self,
        segment_fruits: Vec<<TCollector::Child as SegmentCollector>::Fruit>,
    ) -> crate::Result<TCollector::Fruit> {
        self.collector.merge_fruits(segment_fruits)
    }
}

pub struct FilterSegmentCollector<TSegmentCollector, TPredicate, TPredicateValue>
where
    TPredicate: 'static,
    TPredicateValue: 'static + FastValue,
{
    fast_field_reader: FastFieldReader<TPredicateValue>,
    segment_collector: TSegmentCollector,
    predicate: &'static TPredicate,
    t_predicate_value: PhantomData<TPredicateValue>,
}

impl<TSegmentCollector, TPredicate, TPredicateValue> SegmentCollector
    for FilterSegmentCollector<TSegmentCollector, TPredicate, TPredicateValue>
where
    TSegmentCollector: SegmentCollector,
    TPredicate: 'static + Fn(TPredicateValue) -> bool + Send + Sync,
    TPredicateValue: 'static + FastValue,
{
    type Fruit = TSegmentCollector::Fruit;

    fn collect(&mut self, doc: u32, score: Score) {
        let value = self.fast_field_reader.get(doc);
        if (self.predicate)(value) {
            self.segment_collector.collect(doc, score)
        }
    }

    fn harvest(self) -> <TSegmentCollector as SegmentCollector>::Fruit {
        self.segment_collector.harvest()
    }
}
<commit_msg>change unit test<commit_after>// # Custom collector example
//
// This example shows how you can implement your own
// collector. As an example, we will compute a collector
// that computes the standard deviation of a given fast field.
//
// Of course, you can have a look at the tantivy's built-in collectors
// such as the `CountCollector` for more examples.

// ---
// Importing tantivy...
use std::marker::PhantomData;

use crate::collector::{Collector, SegmentCollector};
use crate::fastfield::{FastFieldReader, FastValue};
use crate::schema::Field;
use crate::{Score, SegmentReader, TantivyError};

/// The `FilterCollector` collector filters docs using a u64 fast field value and a predicate.
/// Only the documents for which the predicate returned "true" will be passed on to the next collector.
///
/// ```rust
/// use tantivy::collector::{TopDocs, FilterCollector};
/// use tantivy::query::QueryParser;
/// use tantivy::schema::{Schema, TEXT, INDEXED, FAST};
/// use tantivy::{doc, DocAddress, Index};
///
/// let mut schema_builder = Schema::builder();
/// let title = schema_builder.add_text_field("title", TEXT);
/// let price = schema_builder.add_u64_field("price", INDEXED | FAST);
/// let schema = schema_builder.build();
/// let index = Index::create_in_ram(schema);
///
/// let mut index_writer = index.writer_with_num_threads(1, 10_000_000).unwrap();
/// index_writer.add_document(doc!(title => "The Name of the Wind", price => 30_200u64));
/// index_writer.add_document(doc!(title => "The Diary of Muadib", price => 29_240u64));
/// index_writer.add_document(doc!(title => "A Dairy Cow", price => 21_240u64));
/// index_writer.add_document(doc!(title => "The Diary of a Young Girl", price => 20_120u64));
/// assert!(index_writer.commit().is_ok());
///
/// let reader = index.reader().unwrap();
/// let searcher = reader.searcher();
///
/// let query_parser = QueryParser::for_index(&index, vec![title]);
/// let query = query_parser.parse_query("diary").unwrap();
/// let no_filter_collector = FilterCollector::new(price, &|value: u64| value > 20_120u64, TopDocs::with_limit(2));
/// let top_docs = searcher.search(&query, &no_filter_collector).unwrap();
///
/// assert_eq!(top_docs.len(), 1);
/// assert_eq!(top_docs[0].1, DocAddress(0, 1));
///
/// let filter_all_collector: FilterCollector<_, _, u64> = FilterCollector::new(price, &|value| value < 5u64, TopDocs::with_limit(2));
/// let filtered_top_docs = searcher.search(&query, &filter_all_collector).unwrap();
///
/// assert_eq!(filtered_top_docs.len(), 0);
/// ```
pub struct FilterCollector<TCollector, TPredicate, TPredicateValue: FastValue>
where
    TPredicate: 'static,
{
    field: Field,
    collector: TCollector,
    predicate: &'static TPredicate,
    t_predicate_value: PhantomData<TPredicateValue>,
}

impl<TCollector, TPredicate, TPredicateValue: FastValue>
    FilterCollector<TCollector, TPredicate, TPredicateValue>
where
    TCollector: Collector + Send + Sync,
    TPredicate: Fn(TPredicateValue) -> bool + Send + Sync,
{
    /// Create a new FilterCollector.
    pub fn new(
        field: Field,
        predicate: &'static TPredicate,
        collector: TCollector,
    ) -> FilterCollector<TCollector, TPredicate, TPredicateValue> {
        FilterCollector {
            field,
            predicate,
            collector,
            t_predicate_value: PhantomData,
        }
    }
}

impl<TCollector, TPredicate, TPredicateValue: FastValue> Collector
    for FilterCollector<TCollector, TPredicate, TPredicateValue>
where
    TCollector: Collector + Send + Sync,
    TPredicate: 'static + Fn(TPredicateValue) -> bool + Send + Sync,
    TPredicateValue: 'static + FastValue,
{
    // That's the type of our result.
    // Our standard deviation will be a float.
    type Fruit = TCollector::Fruit;

    type Child = FilterSegmentCollector<TCollector::Child, TPredicate, TPredicateValue>;

    fn for_segment(
        &self,
        segment_local_id: u32,
        segment_reader: &SegmentReader,
    ) -> crate::Result<FilterSegmentCollector<TCollector::Child, TPredicate, TPredicateValue>> {
        let schema = segment_reader.schema();
        let field_entry = schema.get_field_entry(self.field);
        if !field_entry.is_fast() {
            return Err(TantivyError::SchemaError(format!(
                "Field {:?} is not a fast field.",
                field_entry.name()
            )));
        }
        let requested_type = TPredicateValue::to_type();
        let field_schema_type = field_entry.field_type().value_type();
        if requested_type != field_schema_type {
            return Err(TantivyError::SchemaError(format!(
                "Field {:?} is of type {:?}!={:?}",
                field_entry.name(),
                requested_type,
                field_schema_type
            )));
        }

        let fast_field_reader = segment_reader
            .fast_fields()
            .typed_fast_field_reader(self.field)
            .ok_or_else(|| {
                TantivyError::SchemaError(format!(
                    "{:?} is not declared as a fast field in the schema.",
                    self.field
                ))
            })?;

        let segment_collector = self
            .collector
            .for_segment(segment_local_id, segment_reader)?;

        Ok(FilterSegmentCollector {
            fast_field_reader,
            segment_collector: segment_collector,
            predicate: self.predicate,
            t_predicate_value: PhantomData,
        })
    }

    fn requires_scoring(&self) -> bool {
        self.collector.requires_scoring()
    }

    fn merge_fruits(
        &self,
        segment_fruits: Vec<<TCollector::Child as SegmentCollector>::Fruit>,
    ) -> crate::Result<TCollector::Fruit> {
        self.collector.merge_fruits(segment_fruits)
    }
}

pub struct FilterSegmentCollector<TSegmentCollector, TPredicate, TPredicateValue>
where
    TPredicate: 'static,
    TPredicateValue: 'static + FastValue,
{
    fast_field_reader: FastFieldReader<TPredicateValue>,
    segment_collector: TSegmentCollector,
    predicate: &'static TPredicate,
    t_predicate_value: PhantomData<TPredicateValue>,
}

impl<TSegmentCollector, TPredicate, TPredicateValue> SegmentCollector
    for FilterSegmentCollector<TSegmentCollector, TPredicate, TPredicateValue>
where
    TSegmentCollector: SegmentCollector,
    TPredicate: 'static + Fn(TPredicateValue) -> bool + Send + Sync,
    TPredicateValue: 'static + FastValue,
{
    type Fruit = TSegmentCollector::Fruit;

    fn collect(&mut self, doc: u32, score: Score) {
        let value = self.fast_field_reader.get(doc);
        if (self.predicate)(value) {
            self.segment_collector.collect(doc, score)
        }
    }

    fn harvest(self) -> <TSegmentCollector as SegmentCollector>::Fruit {
        self.segment_collector.harvest()
    }
}
<|endoftext|>


<commit_before>use std::cmp;

use rand;

use errors::{Result, ErrorKind};

/// Overhead per chunk is 12 bytes: magic(2) + id(8) + pos(1) + total (1)
const CHUNK_OVERHEAD: u8 = 12;

/// Chunk-size for LANs
const CHUNK_SIZE_LAN: u16 = 8154;

/// Chunk-size for WANs
const CHUNK_SIZE_WAN: u16 = 1420;

/// Magic bytes identifying a GELF message chunk
static MAGIC_BYTES: &'static [u8; 2] = b"\x1e\x0f";

/// ChunkSize is a value type representing the size of a message-chunk
///
/// It provides default sizes for WANs and LANs
#[derive(Clone, Copy, Debug)]
pub enum ChunkSize {
    LAN,
    WAN,
    Custom(u16),
}

impl ChunkSize {
    /// Return the size associated with the chunk-size
    pub fn size(&self) -> u16 {
        match *self {
            ChunkSize::LAN => CHUNK_SIZE_LAN,
            ChunkSize::WAN => CHUNK_SIZE_WAN,
            ChunkSize::Custom(size) => size,
        }
    }
}

/// ChunkedMessage is an internal type for chunking an already serialized `WireMessage`
pub struct ChunkedMessage {
    chunk_size: ChunkSize,
    payload: Vec<u8>,
    num_chunks: u8,
    id: ChunkedMessageId,
}

impl ChunkedMessage {
    /// Construct a new ChunkedMessage
    ///
    /// Several sanity checks are performed on construction:
    /// - chunk_size must be greater than 0
    /// - GELF allows for a maximum of 128 chunks per message
    pub fn new(chunk_size: ChunkSize, message: Vec<u8>) -> Result<ChunkedMessage> {

        if chunk_size.size() == 0 {
            bail!(ErrorKind::IllegalChunkSize(chunk_size.size()));
        }

        // Ceiled integer division with (a + b - 1) / b
        // Calculate with 64bit integers to avoid overflow - maybe replace with checked_*?
        let size = chunk_size.size() as u64;
        let num_chunks = (message.len() as u64 + size as u64 - 1) / size;

        if num_chunks > 128 {
            bail!(ErrorKind::ChunkMessageFailed("Number of chunks exceeds 128, which the the \
                                                 maximum number of chunks in GELF. Check your \
                                                 chunk_size"))
        }

        Ok(ChunkedMessage {
            chunk_size: chunk_size,
            payload: message,
            id: ChunkedMessageId::random(),
            num_chunks: num_chunks as u8,
        })
    }

    /// Return the byte-length of the chunked message inclduing all overhead
    pub fn len(&self) -> u64 {
        if self.num_chunks > 1 {
            self.payload.len() as u64 + self.num_chunks as u64 * CHUNK_OVERHEAD as u64
        } else {
            self.payload.len() as u64
        }
    }

    /// Return an iterator over all chunks of the message
    pub fn iter(&self) -> ChunkedMessageIterator {
        ChunkedMessageIterator::new(self)
    }
}

/// An iterator over all a chunked message's chunks
///
/// This always be constructed by `ChunkedMessage`
pub struct ChunkedMessageIterator<'a> {
    chunk_num: u8,
    message: &'a ChunkedMessage,
}

impl<'a> ChunkedMessageIterator<'a> {
    /// Create a new ChunkedMessageIterator
    fn new(msg: &'a ChunkedMessage) -> ChunkedMessageIterator {
        ChunkedMessageIterator {
            message: msg,
            chunk_num: 0,
        }
    }
}

impl<'a> Iterator for ChunkedMessageIterator<'a> {
    type Item = Vec<u8>;

    /// Returns the next chunk (if existant)
    fn next(&mut self) -> Option<Vec<u8>> {
        if self.chunk_num >= self.message.num_chunks {
            return None;
        }

        let mut chunk = Vec::new();

        // Set the chunks boundaries
        let chunk_size = self.message.chunk_size.size();
        let slice_start = (self.chunk_num as u16 * chunk_size) as usize;
        let slice_end = cmp::min(slice_start + chunk_size as usize,
                                 self.message.payload.len());

        // The chunk header is only required when the message size exceeds one chunk
        if self.message.num_chunks > 1 {
            // Chunk binary layout:
            //  2 bytes (magic bytes)
            //  8 bytes (message id)
            //  1 byte  (chunk number)
            //  1 byte  (total amount of chunks in this message)
            //  n bytes (chunk payload)
            chunk.extend(MAGIC_BYTES.iter());
            chunk.extend(self.message.id.as_bytes());
            chunk.push(self.chunk_num);
            chunk.push(self.message.num_chunks);
        }

        chunk.extend(self.message.payload[slice_start..slice_end].iter());

        self.chunk_num += 1;

        Some(chunk)
    }
}

/// The representation of a chunked message id
///
/// Every chunked message requires an ID which consists of 8 bytes. This is the same
/// as an 64bit integer. This struct provides some convenience functions on this type.
struct ChunkedMessageId([u8; 8]);

#[allow(dead_code)]
impl<'a> ChunkedMessageId {
    /// Create a new, random ChunkedMessageId.
    fn random() -> ChunkedMessageId {
        let mut bytes = [0; 8];

        for b in 0..8 {
            bytes[b] = rand::random();
        }

        return ChunkedMessageId::from_bytes(bytes);
    }

    /// Create a new ChunkedMessageId from a 64 int.
    fn from_int(mut id: u64) -> ChunkedMessageId {
        let mut bytes = [0; 8];
        for i in 0..8 {
            bytes[7 - i] = (id & 0xff) as u8;
            id >>= 8;
        }

        ChunkedMessageId(bytes)
    }

    /// Create a new ChunkedMessageId from a byte-array.
    fn from_bytes(bytes: [u8; 8]) -> ChunkedMessageId {
        ChunkedMessageId(bytes)
    }

    /// Return the message id as a byte-slice.
    fn as_bytes(&self) -> &[u8; 8] {
        &self.0
    }

    /// Return the message id as an 64bit uint.
    fn to_int(&self) -> u64 {
        self.0.iter().fold(0_u64, |id, &i| id << 8 | i as u64)
    }
}

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

    #[test]
    fn chunked_message_id_from_and_to_bytes() {
        let raw_ids = vec![b"\xff\xff\xff\xff\xff\xff\xff\xff",
                           b"\x00\x00\x00\x00\x00\x00\x00\x00",
                           b"\xaa\xaa\xaa\xaa\xaa\xaa\xaa\xaa",
                           b"\x55\x55\x55\x55\x55\x55\x55\x55",
                           b"\x00\x01\x02\x03\x04\x05\x06\x07",
                           b"\x07\x06\x05\x04\x03\x02\x01\x00",
                           b"\x00\x10\x20\x30\x40\x50\x60\x70",
                           b"\x70\x60\x50\x40\x30\x20\x10\x00"];

        for raw_id in raw_ids {
            let id = ChunkedMessageId::from_bytes(raw_id.clone());
            assert_eq!(id.as_bytes(), raw_id);
        }
    }

    #[test]
    fn chunked_message_id_from_and_to_int() {
        let raw_ids = vec![0xffffffffffffffff,
                           0x0000000000000000,
                           0xaaaaaaaaaaaaaaaa,
                           0x5555555555555555,
                           0x0001020304050607,
                           0x0706050403020100,
                           0x0010203040506070,
                           0x7060504030201000];

        for raw_id in raw_ids {
            let id = ChunkedMessageId::from_int(raw_id);
            assert_eq!(id.to_int(), raw_id);
        }
    }

    #[test]
    #[should_panic(expected = "Number of chunks")]
    fn fail_too_many_chunks() {
        ChunkedMessage::new(ChunkSize::Custom(1), get_data(129)).unwrap();
    }

    #[test]
    fn chunk_message_len() {
        let msg_1_chunk = ChunkedMessage::new(ChunkSize::Custom(1), get_data(1)).unwrap();
        let msg_2_chunks = ChunkedMessage::new(ChunkSize::Custom(1), get_data(2)).unwrap();
        let msg_128_chunks = ChunkedMessage::new(ChunkSize::Custom(1), get_data(128)).unwrap();

        assert_eq!(msg_1_chunk.len(), 1);
        assert_eq!(msg_2_chunks.len() as u32, 2 + 2 * CHUNK_OVERHEAD as u32);
        assert_eq!(msg_128_chunks.len() as u64,
                   128 + 128 * (CHUNK_OVERHEAD as u64));
    }

    #[test]
    fn chunk_message_id_random() {
        let msg1 = ChunkedMessage::new(ChunkSize::Custom(1), get_data(1)).unwrap();
        let msg2 = ChunkedMessage::new(ChunkSize::Custom(1), get_data(1)).unwrap();
        let msg3 = ChunkedMessage::new(ChunkSize::Custom(1), get_data(1)).unwrap();

        assert!(msg1.id.to_int() != msg2.id.to_int());
        assert!(msg3.id.to_int() != msg2.id.to_int());
        assert!(msg1.id.to_int() != msg3.id.to_int());
    }

    #[test]
    fn chunk_message_correct_math() {
        let msg1 = ChunkedMessage::new(ChunkSize::Custom(3), get_data(1)).unwrap();
        let msg2 = ChunkedMessage::new(ChunkSize::Custom(3), get_data(2)).unwrap();
        let msg3 = ChunkedMessage::new(ChunkSize::Custom(3), get_data(3)).unwrap();
        let msg4 = ChunkedMessage::new(ChunkSize::Custom(3), get_data(4)).unwrap();
        let msg5 = ChunkedMessage::new(ChunkSize::Custom(3), get_data(5)).unwrap();
        let msg6 = ChunkedMessage::new(ChunkSize::Custom(3), get_data(6)).unwrap();
        let msg7 = ChunkedMessage::new(ChunkSize::Custom(3), get_data(7)).unwrap();

        assert_eq!(msg1.num_chunks, 1);
        assert_eq!(msg2.num_chunks, 1);
        assert_eq!(msg3.num_chunks, 1);
        assert_eq!(msg4.num_chunks, 2);
        assert_eq!(msg5.num_chunks, 2);
        assert_eq!(msg6.num_chunks, 2);
        assert_eq!(msg7.num_chunks, 3);
    }

    #[test]
    fn chunk_message_chunking() {
        // 10 exact chunks
        check_chunks(10, 100, 10);

        // 4 inexact chunks
        check_chunks(33, 100, 4);

        // test no chunks
        let msg = ChunkedMessage::new(ChunkSize::Custom(100), get_data(100)).unwrap();
        let mut iter = msg.iter();
        let chunk = iter.next().unwrap();
        assert_eq!(iter.next(), None);
        assert_eq!(chunk.len(), 100);
        assert_eq!(chunk[0], 0);
        assert_eq!(chunk[99], 99);
    }

    #[test]
    #[should_panic]
    fn test_illegal_chunk_size() {
        ChunkedMessage::new(ChunkSize::Custom(0), get_data(1)).unwrap();
    }

    fn get_data(len: usize) -> Vec<u8> {
        let mut data = Vec::with_capacity(len);
        for i in 0..len {
            data.push(i as u8);
        }

        data
    }

    fn check_chunks(chunk_size: u8, msg_size: u8, expected_chunk_count: u8) {
        let msg = ChunkedMessage::new(ChunkSize::Custom(chunk_size as u16),
                                      get_data(msg_size as usize))
            .unwrap();
        let mut counter: u8 = 0;
        for chunk in msg.iter() {
            println!("{:?}", chunk);

            // length is in budget

            assert!(chunk.len() as u8 <= chunk_size + 12);
            // magic bytes
            assert_eq!(chunk[0], MAGIC_BYTES[0]);
            assert_eq!(chunk[1], MAGIC_BYTES[1]);

            // pos/counter section
            assert_eq!(chunk[10], counter);
            assert_eq!(chunk[11], expected_chunk_count);

            // first and last byte
            assert_eq!(chunk[12], counter * chunk_size);
            assert_eq!(*chunk.last().unwrap(),
                       ::std::cmp::min((counter + 1) * chunk_size, 100) - 1);

            counter += 1;
        }

        assert_eq!(counter, expected_chunk_count);
    }

}
<commit_msg>fix `attempt to multiply with overflow`<commit_after>use std::cmp;

use rand;

use errors::{Result, ErrorKind};

/// Overhead per chunk is 12 bytes: magic(2) + id(8) + pos(1) + total (1)
const CHUNK_OVERHEAD: u8 = 12;

/// Chunk-size for LANs
const CHUNK_SIZE_LAN: u16 = 8154;

/// Chunk-size for WANs
const CHUNK_SIZE_WAN: u16 = 1420;

/// Magic bytes identifying a GELF message chunk
static MAGIC_BYTES: &'static [u8; 2] = b"\x1e\x0f";

/// ChunkSize is a value type representing the size of a message-chunk
///
/// It provides default sizes for WANs and LANs
#[derive(Clone, Copy, Debug)]
pub enum ChunkSize {
    LAN,
    WAN,
    Custom(u16),
}

impl ChunkSize {
    /// Return the size associated with the chunk-size
    pub fn size(&self) -> u16 {
        match *self {
            ChunkSize::LAN => CHUNK_SIZE_LAN,
            ChunkSize::WAN => CHUNK_SIZE_WAN,
            ChunkSize::Custom(size) => size,
        }
    }
}

/// ChunkedMessage is an internal type for chunking an already serialized `WireMessage`
pub struct ChunkedMessage {
    chunk_size: ChunkSize,
    payload: Vec<u8>,
    num_chunks: u8,
    id: ChunkedMessageId,
}

impl ChunkedMessage {
    /// Construct a new ChunkedMessage
    ///
    /// Several sanity checks are performed on construction:
    /// - chunk_size must be greater than 0
    /// - GELF allows for a maximum of 128 chunks per message
    pub fn new(chunk_size: ChunkSize, message: Vec<u8>) -> Result<ChunkedMessage> {

        if chunk_size.size() == 0 {
            bail!(ErrorKind::IllegalChunkSize(chunk_size.size()));
        }

        // Ceiled integer division with (a + b - 1) / b
        // Calculate with 64bit integers to avoid overflow - maybe replace with checked_*?
        let size = chunk_size.size() as u64;
        let num_chunks = (message.len() as u64 + size as u64 - 1) / size;

        if num_chunks > 128 {
            bail!(ErrorKind::ChunkMessageFailed("Number of chunks exceeds 128, which the the \
                                                 maximum number of chunks in GELF. Check your \
                                                 chunk_size"))
        }

        Ok(ChunkedMessage {
            chunk_size: chunk_size,
            payload: message,
            id: ChunkedMessageId::random(),
            num_chunks: num_chunks as u8,
        })
    }

    /// Return the byte-length of the chunked message inclduing all overhead
    pub fn len(&self) -> u64 {
        if self.num_chunks > 1 {
            self.payload.len() as u64 + self.num_chunks as u64 * CHUNK_OVERHEAD as u64
        } else {
            self.payload.len() as u64
        }
    }

    /// Return an iterator over all chunks of the message
    pub fn iter(&self) -> ChunkedMessageIterator {
        ChunkedMessageIterator::new(self)
    }
}

/// An iterator over all a chunked message's chunks
///
/// This always be constructed by `ChunkedMessage`
pub struct ChunkedMessageIterator<'a> {
    chunk_num: u8,
    message: &'a ChunkedMessage,
}

impl<'a> ChunkedMessageIterator<'a> {
    /// Create a new ChunkedMessageIterator
    fn new(msg: &'a ChunkedMessage) -> ChunkedMessageIterator {
        ChunkedMessageIterator {
            message: msg,
            chunk_num: 0,
        }
    }
}

impl<'a> Iterator for ChunkedMessageIterator<'a> {
    type Item = Vec<u8>;

    /// Returns the next chunk (if existant)
    fn next(&mut self) -> Option<Vec<u8>> {
        if self.chunk_num >= self.message.num_chunks {
            return None;
        }

        let mut chunk = Vec::new();

        // Set the chunks boundaries
        let chunk_size = self.message.chunk_size.size();
        let slice_start = (self.chunk_num as u32 * chunk_size as u32) as usize;
        let slice_end = cmp::min(slice_start + chunk_size as usize,
                                 self.message.payload.len());

        // The chunk header is only required when the message size exceeds one chunk
        if self.message.num_chunks > 1 {
            // Chunk binary layout:
            //  2 bytes (magic bytes)
            //  8 bytes (message id)
            //  1 byte  (chunk number)
            //  1 byte  (total amount of chunks in this message)
            //  n bytes (chunk payload)
            chunk.extend(MAGIC_BYTES.iter());
            chunk.extend(self.message.id.as_bytes());
            chunk.push(self.chunk_num);
            chunk.push(self.message.num_chunks);
        }

        chunk.extend(self.message.payload[slice_start..slice_end].iter());

        self.chunk_num += 1;

        Some(chunk)
    }
}

/// The representation of a chunked message id
///
/// Every chunked message requires an ID which consists of 8 bytes. This is the same
/// as an 64bit integer. This struct provides some convenience functions on this type.
struct ChunkedMessageId([u8; 8]);

#[allow(dead_code)]
impl<'a> ChunkedMessageId {
    /// Create a new, random ChunkedMessageId.
    fn random() -> ChunkedMessageId {
        let mut bytes = [0; 8];

        for b in 0..8 {
            bytes[b] = rand::random();
        }

        return ChunkedMessageId::from_bytes(bytes);
    }

    /// Create a new ChunkedMessageId from a 64 int.
    fn from_int(mut id: u64) -> ChunkedMessageId {
        let mut bytes = [0; 8];
        for i in 0..8 {
            bytes[7 - i] = (id & 0xff) as u8;
            id >>= 8;
        }

        ChunkedMessageId(bytes)
    }

    /// Create a new ChunkedMessageId from a byte-array.
    fn from_bytes(bytes: [u8; 8]) -> ChunkedMessageId {
        ChunkedMessageId(bytes)
    }

    /// Return the message id as a byte-slice.
    fn as_bytes(&self) -> &[u8; 8] {
        &self.0
    }

    /// Return the message id as an 64bit uint.
    fn to_int(&self) -> u64 {
        self.0.iter().fold(0_u64, |id, &i| id << 8 | i as u64)
    }
}

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

    #[test]
    fn chunked_message_id_from_and_to_bytes() {
        let raw_ids = vec![b"\xff\xff\xff\xff\xff\xff\xff\xff",
                           b"\x00\x00\x00\x00\x00\x00\x00\x00",
                           b"\xaa\xaa\xaa\xaa\xaa\xaa\xaa\xaa",
                           b"\x55\x55\x55\x55\x55\x55\x55\x55",
                           b"\x00\x01\x02\x03\x04\x05\x06\x07",
                           b"\x07\x06\x05\x04\x03\x02\x01\x00",
                           b"\x00\x10\x20\x30\x40\x50\x60\x70",
                           b"\x70\x60\x50\x40\x30\x20\x10\x00"];

        for raw_id in raw_ids {
            let id = ChunkedMessageId::from_bytes(raw_id.clone());
            assert_eq!(id.as_bytes(), raw_id);
        }
    }

    #[test]
    fn chunked_message_id_from_and_to_int() {
        let raw_ids = vec![0xffffffffffffffff,
                           0x0000000000000000,
                           0xaaaaaaaaaaaaaaaa,
                           0x5555555555555555,
                           0x0001020304050607,
                           0x0706050403020100,
                           0x0010203040506070,
                           0x7060504030201000];

        for raw_id in raw_ids {
            let id = ChunkedMessageId::from_int(raw_id);
            assert_eq!(id.to_int(), raw_id);
        }
    }

    #[test]
    #[should_panic(expected = "Number of chunks")]
    fn fail_too_many_chunks() {
        ChunkedMessage::new(ChunkSize::Custom(1), get_data(129)).unwrap();
    }

    #[test]
    fn chunk_message_len() {
        let msg_1_chunk = ChunkedMessage::new(ChunkSize::Custom(1), get_data(1)).unwrap();
        let msg_2_chunks = ChunkedMessage::new(ChunkSize::Custom(1), get_data(2)).unwrap();
        let msg_128_chunks = ChunkedMessage::new(ChunkSize::Custom(1), get_data(128)).unwrap();

        assert_eq!(msg_1_chunk.len(), 1);
        assert_eq!(msg_2_chunks.len() as u32, 2 + 2 * CHUNK_OVERHEAD as u32);
        assert_eq!(msg_128_chunks.len() as u64,
                   128 + 128 * (CHUNK_OVERHEAD as u64));
    }

    #[test]
    fn chunk_message_id_random() {
        let msg1 = ChunkedMessage::new(ChunkSize::Custom(1), get_data(1)).unwrap();
        let msg2 = ChunkedMessage::new(ChunkSize::Custom(1), get_data(1)).unwrap();
        let msg3 = ChunkedMessage::new(ChunkSize::Custom(1), get_data(1)).unwrap();

        assert!(msg1.id.to_int() != msg2.id.to_int());
        assert!(msg3.id.to_int() != msg2.id.to_int());
        assert!(msg1.id.to_int() != msg3.id.to_int());
    }

    #[test]
    fn chunk_message_correct_math() {
        let msg1 = ChunkedMessage::new(ChunkSize::Custom(3), get_data(1)).unwrap();
        let msg2 = ChunkedMessage::new(ChunkSize::Custom(3), get_data(2)).unwrap();
        let msg3 = ChunkedMessage::new(ChunkSize::Custom(3), get_data(3)).unwrap();
        let msg4 = ChunkedMessage::new(ChunkSize::Custom(3), get_data(4)).unwrap();
        let msg5 = ChunkedMessage::new(ChunkSize::Custom(3), get_data(5)).unwrap();
        let msg6 = ChunkedMessage::new(ChunkSize::Custom(3), get_data(6)).unwrap();
        let msg7 = ChunkedMessage::new(ChunkSize::Custom(3), get_data(7)).unwrap();

        assert_eq!(msg1.num_chunks, 1);
        assert_eq!(msg2.num_chunks, 1);
        assert_eq!(msg3.num_chunks, 1);
        assert_eq!(msg4.num_chunks, 2);
        assert_eq!(msg5.num_chunks, 2);
        assert_eq!(msg6.num_chunks, 2);
        assert_eq!(msg7.num_chunks, 3);
    }

    #[test]
    fn chunk_message_chunking() {
        // 10 exact chunks
        check_chunks(10, 100, 10);

        // 4 inexact chunks
        check_chunks(33, 100, 4);

        // test no chunks
        let msg = ChunkedMessage::new(ChunkSize::Custom(100), get_data(100)).unwrap();
        let mut iter = msg.iter();
        let chunk = iter.next().unwrap();
        assert_eq!(iter.next(), None);
        assert_eq!(chunk.len(), 100);
        assert_eq!(chunk[0], 0);
        assert_eq!(chunk[99], 99);
    }

    #[test]
    fn chunk_large_message_chunking() {
        // 100k of msg
        chunking(CHUNK_SIZE_WAN, 100000);
    }

    fn chunking(chunk_size: u16, msg_size: u32) {
        check_chunks(chunk_size as u16, msg_size, (msg_size / chunk_size as u32) as u8 + 1);
    }

    #[test]
    #[should_panic]
    fn test_illegal_chunk_size() {
        ChunkedMessage::new(ChunkSize::Custom(0), get_data(1)).unwrap();
    }

    fn get_data(len: usize) -> Vec<u8> {
        let mut data = Vec::with_capacity(len);
        for i in 0..len {
            data.push(i as u8);
        }

        data
    }

    fn check_chunks(chunk_size: u16, msg_size: u32, expected_chunk_count: u8) {
        let msg_data = get_data(msg_size as usize);
        let msg_data_clone = msg_data.clone();
        let msg = ChunkedMessage::new(ChunkSize::Custom(chunk_size as u16),
                                      msg_data)
            .unwrap();
        let mut counter: u8 = 0;
        for chunk in msg.iter() {
            println!("{:?}", chunk);

            // length is in budget

            assert!(chunk.len() as u16 <= chunk_size + 12);
            // magic bytes
            assert_eq!(chunk[0], MAGIC_BYTES[0]);
            assert_eq!(chunk[1], MAGIC_BYTES[1]);

            // pos/counter section
            assert_eq!(chunk[10], counter);
            assert_eq!(chunk[11], expected_chunk_count);

            // first and last byte
            let first_index = (counter as u32 * chunk_size as u32) as usize;
            let last_index = (::std::cmp::min((counter as u32 + 1) * chunk_size as u32, msg_size) - 1) as usize;
            assert_eq!(chunk[12], msg_data_clone[first_index]);
            assert_eq!(*chunk.last().unwrap(),
                       msg_data_clone[last_index]);

            counter += 1;
        }

        assert_eq!(counter, expected_chunk_count);
    }

}
<|endoftext|>


<commit_before>use cripes::util::iter::*;

/// Test the `Once` iterator.
#[test]
fn test_once() {
    let mut n = 0;

    let y = "foo";
    for something in Once::new(y) { n += 1; println!("{:?}", something); }
    assert_eq!(1, n);

    let things = vec!['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z'];
    n = 0;
    for something in Once::new(&things) { n += 1; println!("{:?}", something); }
    assert_eq!(1, n);

    n = 0;
    for something in things.iter() {
        for _ in Once::new(something) { n += 1 }
    }
    assert_eq!(n, 26);
}

/// test for the `Empty` iterator.
#[test]
fn test_empty() {
    let mut n = 0;
    for _ in Empty::<()>::new() { n += 1; }
    assert_eq!(n, 0);
}
<commit_msg>test_iter.rs: add an extra simple test for Once<commit_after>use cripes::util::iter::*;

/// Test the `Once` iterator.
#[test]
fn test_once() {
    assert_eq!(1, Once::new(42).count());

    let mut n = 0;
    let y = "foo";
    for something in Once::new(y) { n += 1; println!("{:?}", something); }
    assert_eq!(1, n);

    let things = vec!['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z'];
    n = 0;
    for something in Once::new(&things) { n += 1; println!("{:?}", something); }
    assert_eq!(1, n);

    n = 0;
    for something in things.iter() {
        for _ in Once::new(something) { n += 1 }
    }
    assert_eq!(n, 26);
}

/// test for the `Empty` iterator.
#[test]
fn test_empty() {
    let mut n = 0;
    for _ in Empty::<()>::new() { n += 1; }
    assert_eq!(n, 0);
}
<|endoftext|>