<commit_before>use std::fmt::{self, Debug, Display};
use std::mem;

use nom::{self, le_u32};
use quick_error::ResultExt;

// Generated with ./mk_vsl_tag from Varnish headers: include/tbl/vsl_tags.h include/tbl/vsl_tags_http.h include/vsl_int.h
// https://github.com/varnishcache/varnish-cache/blob/master/include/vapi/vsl_int.h
// https://github.com/varnishcache/varnish-cache/blob/master/include/tbl/vsl_tags.h
// https://github.com/varnishcache/varnish-cache/blob/master/include/tbl/vsl_tags_http.h
mod tag_e;
pub use self::tag_e::VSL_tag_e as VslRecordTag;

mod message_parsers;
pub use self::message_parsers::*;

mod maybe_string;
pub use self::maybe_string::{MaybeStr, MaybeString};

mod store;
pub use self::store::VslStore;

/*
 * Shared memory log format
 *
 * The log member points to an array of 32bit unsigned integers containing
 * log records.
 *
 * Each logrecord consist of:
 *    [n]               = ((type & 0xff) << 24) | (length & 0xffff)
 *    [n + 1]           = ((marker & 0x03) << 30) | (identifier & 0x3fffffff)
 *    [n + 2] ... [m]   = content (NUL-terminated)
 *
 * Logrecords are NUL-terminated so that string functions can be run
 * directly on the shmlog data.
 *
 * Notice that the constants in these macros cannot be changed without
 * changing corresponding magic numbers in varnishd/cache/cache_shmlog.c
 */

const VSL_LENOFFSET: u32 = 24;
const VSL_LENMASK: u32 = 0xffff;
const VSL_IDENTOFFSET: u8 = 30;
const VSL_IDENTMASK: u32 = !(0b0000_0011 << VSL_IDENTOFFSET);

/*
* VSL_CLIENT(ptr)
*   Non-zero if this is a client transaction
*
* VSL_BACKEND(ptr)
*   Non-zero if this is a backend transaction
*/

bitflags! {
    pub flags Marker: u8 {
        const VSL_CLIENTMARKER  = 0b0000_0001,
        const VSL_BACKENDMARKER = 0b0000_0010,
    }
}

impl Display for Marker {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "[{}{}]",
               if self.contains(VSL_CLIENTMARKER) { "C" } else { " " },
               if self.contains(VSL_BACKENDMARKER) { "B" } else { " " })
    }
}

pub type VslIdent = u32;

named!(pub binary_vsl_tag<&[u8], Option<&[u8]> >, opt!(complete!(tag!(b"VSL\0"))));

#[derive(Debug)]
struct VslRecordHeader {
    tag: u8,
    len: u16,
    marker: Marker,
    ident: VslIdent,
}

fn vsl_record_header<'b>(input: &'b[u8]) -> nom::IResult<&'b[u8], VslRecordHeader, u32> {
    chain!(
        input, r1: le_u32 ~ r2: le_u32,
        || VslRecordHeader {
            tag: (r1 >> VSL_LENOFFSET) as u8,
            len: (r1 & VSL_LENMASK) as u16,
            marker: Marker::from_bits_truncate(((r2 & !VSL_IDENTMASK) >> VSL_IDENTOFFSET) as u8),
            ident: r2 & VSL_IDENTMASK,
        })
}

pub struct VslRecord<'b> {
    pub tag: VslRecordTag,
    pub marker: Marker,
    pub ident: VslIdent,
    pub data: &'b[u8],
}

quick_error! {
    #[derive(Debug)]
    pub enum VslRecordParseError {
        Nom(nom_err: String, tag: VslRecordTag, record: String) {
            context(record: &'a VslRecord<'a>, err: nom::Err<&'a [u8]>) -> (format!("{}", err), record.tag, format!("{}", record))
            display("Nom parser failed on {}: {}", record, nom_err)
        }
    }
}

impl<'b> VslRecord<'b> {
    pub fn parse_data<T, P>(&'b self, parser: P) -> Result<T, VslRecordParseError> where
    P: Fn(&'b [u8]) -> nom::IResult<&'b [u8], T> {
        // Note: need type annotaion for the u32 error type as the output IResult has no Error
        // variant that would help to infer it
        let result: nom::IResult<_, Result<T, _>, u32> = opt_res!(self.data, complete!(parser));
        // unwrap here is safe as complete! eliminates Incomplete variant and opt_res! remaining Error variant
        result.unwrap().1.context(self).map_err(|err| From::from(err))
    }

    pub fn is_client(&self) -> bool {
        self.marker.contains(VSL_CLIENTMARKER)
    }

    pub fn is_backend(&self) -> bool {
        self.marker.contains(VSL_BACKENDMARKER)
    }

    #[cfg(test)]
    pub fn from_str<'s>(tag: VslRecordTag, ident: VslIdent, message: &'s str) -> VslRecord<'s> {
        VslRecord {
            tag: tag,
            marker: VSL_CLIENTMARKER,
            ident: ident,
            data: message.as_ref()
        }
    }
}

impl<'b> Debug for VslRecord<'b> {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        f.debug_struct("VSL Record")
            .field("tag", &self.tag)
            .field("marker", &self.marker)
            .field("ident", &self.ident)
            .field("data", &MaybeStr::from_bytes(&self.data))
            .finish()
    }
}

impl<'b> Display for VslRecord<'b> {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        let tag = format!("{:?}", self.tag);

        if f.alternate() {
            write!(f, "{} {:5} {:18} {}", self.marker, self.ident, tag, MaybeStr::from_bytes(self.data))
        } else {
            write!(f, "VSL record (marker: {} ident: {} tag: {} data: {:?})", self.marker, self.ident, tag, MaybeStr::from_bytes(self.data))
        }
    }
}

fn to_vsl_record_tag(num: u8) -> VslRecordTag {
    // Tend to work even for missing tags as they end up as SLT__Bogus (0)
    unsafe { mem::transmute(num as u32) }
}

pub fn vsl_record_v4<'b>(input: &'b[u8]) -> nom::IResult<&'b[u8], VslRecord<'b>, u32> {
    chain!(
        input,
        header: vsl_record_header ~
        data: take!(header.len - 1) ~ take!(1) ~ take!((4 - header.len % 4) % 4),
        || VslRecord {
            tag: to_vsl_record_tag(header.tag),
            marker: header.marker,
            ident: header.ident,
            data: data
        })
}
<commit_msg>wording<commit_after>use std::fmt::{self, Debug, Display};
use std::mem;

use nom::{self, le_u32};
use quick_error::ResultExt;

// Generated with ./mk_vsl_tag from Varnish headers: include/tbl/vsl_tags.h include/tbl/vsl_tags_http.h include/vsl_int.h
// https://github.com/varnishcache/varnish-cache/blob/master/include/vapi/vsl_int.h
// https://github.com/varnishcache/varnish-cache/blob/master/include/tbl/vsl_tags.h
// https://github.com/varnishcache/varnish-cache/blob/master/include/tbl/vsl_tags_http.h
mod tag_e;
pub use self::tag_e::VSL_tag_e as VslRecordTag;

mod message_parsers;
pub use self::message_parsers::*;

mod maybe_string;
pub use self::maybe_string::{MaybeStr, MaybeString};

mod store;
pub use self::store::VslStore;

/*
 * Shared memory log format
 *
 * The log member points to an array of 32bit unsigned integers containing
 * log records.
 *
 * Each logrecord consist of:
 *    [n]               = ((type & 0xff) << 24) | (length & 0xffff)
 *    [n + 1]           = ((marker & 0x03) << 30) | (identifier & 0x3fffffff)
 *    [n + 2] ... [m]   = content (NUL-terminated)
 *
 * Logrecords are NUL-terminated so that string functions can be run
 * directly on the shmlog data.
 *
 * Notice that the constants in these macros cannot be changed without
 * changing corresponding magic numbers in varnishd/cache/cache_shmlog.c
 */

const VSL_LENOFFSET: u32 = 24;
const VSL_LENMASK: u32 = 0xffff;
const VSL_IDENTOFFSET: u8 = 30;
const VSL_IDENTMASK: u32 = !(0b0000_0011 << VSL_IDENTOFFSET);

/*
* VSL_CLIENT(ptr)
*   Non-zero if this is a client transaction
*
* VSL_BACKEND(ptr)
*   Non-zero if this is a backend transaction
*/

bitflags! {
    pub flags Marker: u8 {
        const VSL_CLIENTMARKER  = 0b0000_0001,
        const VSL_BACKENDMARKER = 0b0000_0010,
    }
}

impl Display for Marker {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "[{}{}]",
               if self.contains(VSL_CLIENTMARKER) { "C" } else { " " },
               if self.contains(VSL_BACKENDMARKER) { "B" } else { " " })
    }
}

pub type VslIdent = u32;

named!(pub binary_vsl_tag<&[u8], Option<&[u8]> >, opt!(complete!(tag!(b"VSL\0"))));

#[derive(Debug)]
struct VslRecordHeader {
    tag: u8,
    len: u16,
    marker: Marker,
    ident: VslIdent,
}

fn vsl_record_header<'b>(input: &'b[u8]) -> nom::IResult<&'b[u8], VslRecordHeader, u32> {
    chain!(
        input, r1: le_u32 ~ r2: le_u32,
        || VslRecordHeader {
            tag: (r1 >> VSL_LENOFFSET) as u8,
            len: (r1 & VSL_LENMASK) as u16,
            marker: Marker::from_bits_truncate(((r2 & !VSL_IDENTMASK) >> VSL_IDENTOFFSET) as u8),
            ident: r2 & VSL_IDENTMASK,
        })
}

pub struct VslRecord<'b> {
    pub tag: VslRecordTag,
    pub marker: Marker,
    pub ident: VslIdent,
    pub data: &'b[u8],
}

quick_error! {
    #[derive(Debug)]
    pub enum VslRecordParseError {
        Nom(nom_err: String, tag: VslRecordTag, record: String) {
            context(record: &'a VslRecord<'a>, err: nom::Err<&'a [u8]>) -> (format!("{}", err), record.tag, format!("{}", record))
            display("Nom parser failed on {}: {}", record, nom_err)
        }
    }
}

impl<'b> VslRecord<'b> {
    pub fn parse_data<T, P>(&'b self, parser: P) -> Result<T, VslRecordParseError> where
    P: Fn(&'b [u8]) -> nom::IResult<&'b [u8], T> {
        // Note: need type annotaion for the u32 error type as the output IResult has no Error
        // variant that would help to infer it
        let result: nom::IResult<_, Result<T, _>, u32> = opt_res!(self.data, complete!(parser));
        // unwrap here is safe as complete! eliminates Incomplete variant and opt_res! remaining Error variant
        result.unwrap().1.context(self).map_err(|err| From::from(err))
    }

    pub fn is_client(&self) -> bool {
        self.marker.contains(VSL_CLIENTMARKER)
    }

    pub fn is_backend(&self) -> bool {
        self.marker.contains(VSL_BACKENDMARKER)
    }

    #[cfg(test)]
    pub fn from_str<'s>(tag: VslRecordTag, ident: VslIdent, message: &'s str) -> VslRecord<'s> {
        VslRecord {
            tag: tag,
            marker: VSL_CLIENTMARKER,
            ident: ident,
            data: message.as_ref()
        }
    }
}

impl<'b> Debug for VslRecord<'b> {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        f.debug_struct("VSL Record")
            .field("tag", &self.tag)
            .field("marker", &self.marker)
            .field("ident", &self.ident)
            .field("data", &MaybeStr::from_bytes(&self.data))
            .finish()
    }
}

impl<'b> Display for VslRecord<'b> {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        let tag = format!("{:?}", self.tag);

        if f.alternate() {
            write!(f, "{} {:5} {:18} {}", self.marker, self.ident, tag, MaybeStr::from_bytes(self.data))
        } else {
            write!(f, "VSL record (marker: {} ident: {} tag: {} data: {:?})", self.marker, self.ident, tag, MaybeStr::from_bytes(self.data))
        }
    }
}

fn to_vsl_record_tag(num: u8) -> VslRecordTag {
    // Works even for missing tags as they end up as SLT__Bogus (value 0)
    unsafe { mem::transmute(num as u32) }
}

pub fn vsl_record_v4<'b>(input: &'b[u8]) -> nom::IResult<&'b[u8], VslRecord<'b>, u32> {
    chain!(
        input,
        header: vsl_record_header ~
        data: take!(header.len - 1) ~ take!(1) ~ take!((4 - header.len % 4) % 4),
        || VslRecord {
            tag: to_vsl_record_tag(header.tag),
            marker: header.marker,
            ident: header.ident,
            data: data
        })
}
<|endoftext|>


<commit_before>extern crate walkdir;
extern crate chrono;

use std::path;
use self::walkdir::{WalkDir, DirEntry, Error as WalkDirError};
use std::fs::File;
use std::io::BufReader;
use std::io::BufRead;
use self::chrono::{DateTime, UTC};
use std::error::Error;
use std::str::FromStr;
use std::net::SocketAddrV4;

pub struct ELBLogEntry {
    timestamp: DateTime<UTC>,
    elb_name: String,
    client_address: SocketAddrV4,
    backend_address: SocketAddrV4,
    request_processing_time: f32,
    backend_processing_time: f32,
    response_processing_time: f32,
    elb_status_code: u16,
    backend_status_code: u16,
    received_bytes: u64,
    sent_bytes: u64,
    request_method: String,
    request_url: String,
    request_http_version: String
}

pub fn file_list(dir: &path::Path, filenames: &mut Vec<DirEntry>) -> Result<usize, WalkDirError> {
    for entry in WalkDir::new(dir).min_depth(1) {
        match entry {
            Err(err) => return Err(err),
            Ok(entry) => filenames.push(entry),
        }
    }
    Ok(filenames.len())
}

//TODO Reconsider logging based on the standard interfaces included with Rust.
//TODO We really want to accept a function to handle the parsed lines.
pub fn process_files(runtime_context: &::RuntimeContext, filenames: Vec<walkdir::DirEntry>) -> usize {
    let debug = runtime_context.debug;
    let mut record_count = 0;
    for filename in filenames {
        debug!(debug, "Processing file {}.", filename.path().display());
        match File::open(filename.path()) {
            Ok(file) => {
                let buffered_file = BufReader::new(&file);
                let recs: Vec<_> = buffered_file.lines()
                    .map(|x| {
                        parse_line(&(x.unwrap()))
                    })
                    .collect();
                record_count += recs.len();
                debug!(debug, "Found {} records in file {}.", recs.len(), filename.path().display());
            },
            Err(e) => {
                println!("ERROR: {}", e);
            }
        }
    }

    record_count
}

//TODO Take a look at the error handling once again.  This doesn't feel write given code you've read.
#[derive(Debug)]
pub struct ParsingError{
    property: &'static str,
    inner_description: Box<Error>,
}

#[derive(Debug)]
pub struct ParsingErrors {
    record: String,
    errors: Vec<ParsingError>,
}

const TIMESTAMP: &'static str = "timestamp";
const CLIENT_ADDRESS: &'static str = "client address";
const BACKEND_ADDRESS: &'static str = "backend address";
const REQUEST_PROCESSING_TIME: &'static str = "request processing time";
const BACKEND_PROCESSING_TIME: &'static str = "backend processing time";
const RESPONSE_PROCESSING_TIME: &'static str = "response processing time";
const ELB_STATUS_CODE: &'static str = "ELB status code";
const BE_STATUS_CODE: &'static str = "backend status code";
const RECEIVED_BYTES: &'static str = "received bytes";
const SENT_BYTES: &'static str = "sent bytes";

pub fn parse_line(line: &String) -> Result<Box<ELBLogEntry>, Box<ParsingErrors>> {
    let split_line: Vec<_> = line.split(" ").collect();
    let mut errors: Vec<ParsingError> = Vec::new();

    //TODO consider an enum representing the idices
    let ts = parse_property::<DateTime<UTC>>(split_line[0], TIMESTAMP, &mut errors);
    let clnt_addr = parse_property::<SocketAddrV4>(split_line[2], CLIENT_ADDRESS, &mut errors);
    let be_addr = parse_property::<SocketAddrV4>(split_line[3], BACKEND_ADDRESS, &mut errors);
    let req_proc_time = parse_property::<f32>(split_line[4], REQUEST_PROCESSING_TIME, &mut errors);
    let be_proc_time = parse_property::<f32>(split_line[5], BACKEND_PROCESSING_TIME, &mut errors);
    let res_proc_time = parse_property::<f32>(split_line[6], RESPONSE_PROCESSING_TIME, &mut errors);
    let elb_sc = parse_property::<u16>(split_line[7], ELB_STATUS_CODE, &mut errors);
    let be_sc = parse_property::<u16>(split_line[8], BE_STATUS_CODE, &mut errors);
    let bytes_received = parse_property::<u64>(split_line[9], RECEIVED_BYTES, &mut errors);
    let bytes_sent = parse_property::<u64>(split_line[10], SENT_BYTES, &mut errors);

    if errors.is_empty() {
        Ok(Box::new(
            ELBLogEntry {
                timestamp: ts.unwrap(),
                elb_name: split_line[1].to_string(),
                client_address: clnt_addr.unwrap(),
                backend_address: be_addr.unwrap(),
                request_processing_time: req_proc_time.unwrap(),
                backend_processing_time: be_proc_time.unwrap(),
                response_processing_time: res_proc_time.unwrap(),
                elb_status_code: elb_sc.unwrap(),
                backend_status_code: be_sc.unwrap(),
                received_bytes: bytes_received.unwrap(),
                sent_bytes: bytes_sent.unwrap(),
                request_method: split_line[11].trim_matches('"').to_string(),
                request_url: split_line[12].to_string(),
                request_http_version: split_line[13].trim_matches('"').to_string()
            }
        ))
    } else {
        Err(Box::new(
            ParsingErrors {
                record: line.clone(),
                errors: errors
            }
        ))
    }
}

fn parse_property<T>(raw_prop: &str, prop_name: &'static str, errors: &mut Vec<ParsingError>) -> Option<T>
    where T: FromStr,
    T::Err: Error + 'static,
{
    match raw_prop.parse::<T>() {
        Ok(parsed) => Some(parsed),

        Err(e) => {
            errors.push(
                ParsingError {
                    property: prop_name,
                    inner_description: Box::new(e),
                }
            );
            None
        }
    }
}

//TODO test the error case.

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

    const TEST_LINE: &'static str = "2015-08-15T23:43:05.302180Z elb-name 172.16.1.6:54814 \
    172.16.1.5:9000 0.000039 0.145507 0.00003 200 200 0 7582 \
    \"GET http://some.domain.com:80/path0/path1?param0=p0&param1=p1 HTTP/1.1\" \
    ";

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_http_version() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_http_version, "HTTP/1.1")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_url() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_url, "http://some.domain.com:80/path0/path1?param0=p0&param1=p1")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_method() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_method, "GET")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_sent_bytes() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.sent_bytes, 7582)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_received_bytes() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.received_bytes, 0)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_backend_status_code() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.backend_status_code, 200)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_elb_status_code() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.elb_status_code, 200)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_response_processing_time() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.response_processing_time, 0.00003)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_backend_processing_time() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.backend_processing_time, 0.145507)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_processing_time() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_processing_time, 0.000039)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_backend_address() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.backend_address, "172.16.1.5:9000".parse().unwrap())
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_client_address() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.client_address, "172.16.1.6:54814".parse().unwrap())
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_timestamp() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(format!("{:?}", elb_log_entry.timestamp), "2015-08-15T23:43:05.302180Z")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_elb_name() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.elb_name, "elb-name")
	}
}
<commit_msg>Reworked the error handling again and came up with what I think is a better architecture.<commit_after>extern crate walkdir;
extern crate chrono;

use std::path;
use self::walkdir::{WalkDir, DirEntry, Error as WalkDirError};
use std::fs::File;
use std::io::BufReader;
use std::io::BufRead;
use self::chrono::{DateTime, UTC};
use std::error::Error;
use std::str::FromStr;
use std::net::SocketAddrV4;
use std::num;

#[derive(Debug)]
pub struct ELBLogEntry {
    timestamp: DateTime<UTC>,
    elb_name: String,
    client_address: SocketAddrV4,
    backend_address: SocketAddrV4,
    request_processing_time: f32,
    backend_processing_time: f32,
    response_processing_time: f32,
    elb_status_code: u16,
    backend_status_code: u16,
    received_bytes: u64,
    sent_bytes: u64,
    request_method: String,
    request_url: String,
    request_http_version: String
}

pub fn file_list(dir: &path::Path, filenames: &mut Vec<DirEntry>) -> Result<usize, WalkDirError> {
    for entry in WalkDir::new(dir).min_depth(1) {
        match entry {
            Err(err) => return Err(err),
            Ok(entry) => filenames.push(entry),
        }
    }
    Ok(filenames.len())
}

//TODO Reconsider logging based on the standard interfaces included with Rust.
//TODO We really want to accept a function to handle the parsed lines.
pub fn process_files(runtime_context: &::RuntimeContext, filenames: Vec<walkdir::DirEntry>) -> usize {
    let debug = runtime_context.debug;
    let mut record_count = 0;
    for filename in filenames {
        debug!(debug, "Processing file {}.", filename.path().display());
        match File::open(filename.path()) {
            Ok(file) => {
                let buffered_file = BufReader::new(&file);
                let recs: Vec<_> = buffered_file.lines()
                    .map(|x| {
                        parse_line(&(x.unwrap()))
                    })
                    .collect();
                record_count += recs.len();
                debug!(debug, "Found {} records in file {}.", recs.len(), filename.path().display());
            },
            Err(e) => {
                println!("ERROR: {}", e);
            }
        }
    }

    record_count
}

//TODO Take a look at the error handling once again.  This doesn't feel write given code you've read.

#[derive(Debug)]
pub struct ParsingErrors {
    record: String,
    errors: Vec<ELBLogEntryErrors>,
}

#[derive(Debug, PartialEq)]
enum ELBLogEntryErrors {
    MalformedEntry,
    ParsingError { property: &'static str, description: String },
}

const ELB_LOG_ENTRY_FIELD_COUNT: usize = 14;
const TIMESTAMP: &'static str = "timestamp";
const CLIENT_ADDRESS: &'static str = "client address";
const BACKEND_ADDRESS: &'static str = "backend address";
const REQUEST_PROCESSING_TIME: &'static str = "request processing time";
const BACKEND_PROCESSING_TIME: &'static str = "backend processing time";
const RESPONSE_PROCESSING_TIME: &'static str = "response processing time";
const ELB_STATUS_CODE: &'static str = "ELB status code";
const BE_STATUS_CODE: &'static str = "backend status code";
const RECEIVED_BYTES: &'static str = "received bytes";
const SENT_BYTES: &'static str = "sent bytes";

pub fn parse_line(line: &String) -> Result<Box<ELBLogEntry>, ParsingErrors> {
    let split_line: Vec<_> = line.split(" ").collect();
    let mut errors: Vec<ELBLogEntryErrors> = Vec::new();
    if split_line.len() == ELB_LOG_ENTRY_FIELD_COUNT {
        //TODO consider an enum representing the idices
        let ts = parse_property::<DateTime<UTC>>(split_line[0], TIMESTAMP, &mut errors);
        let clnt_addr = parse_property::<SocketAddrV4>(split_line[2], CLIENT_ADDRESS, &mut errors);
        let be_addr = parse_property::<SocketAddrV4>(split_line[3], BACKEND_ADDRESS, &mut errors);
        let req_proc_time = parse_property::<f32>(split_line[4], REQUEST_PROCESSING_TIME, &mut errors);
        let be_proc_time = parse_property::<f32>(split_line[5], BACKEND_PROCESSING_TIME, &mut errors);
        let res_proc_time = parse_property::<f32>(split_line[6], RESPONSE_PROCESSING_TIME, &mut errors);
        let elb_sc = parse_property::<u16>(split_line[7], ELB_STATUS_CODE, &mut errors);
        let be_sc = parse_property::<u16>(split_line[8], BE_STATUS_CODE, &mut errors);
        let bytes_received = parse_property::<u64>(split_line[9], RECEIVED_BYTES, &mut errors);
        let bytes_sent = parse_property::<u64>(split_line[10], SENT_BYTES, &mut errors);

        if errors.is_empty() {
            Ok(Box::new(
                ELBLogEntry {
                    timestamp: ts.unwrap(),
                    elb_name: split_line[1].to_string(),
                    client_address: clnt_addr.unwrap(),
                    backend_address: be_addr.unwrap(),
                    request_processing_time: req_proc_time.unwrap(),
                    backend_processing_time: be_proc_time.unwrap(),
                    response_processing_time: res_proc_time.unwrap(),
                    elb_status_code: elb_sc.unwrap(),
                    backend_status_code: be_sc.unwrap(),
                    received_bytes: bytes_received.unwrap(),
                    sent_bytes: bytes_sent.unwrap(),
                    request_method: split_line[11].trim_matches('"').to_string(),
                    request_url: split_line[12].to_string(),
                    request_http_version: split_line[13].trim_matches('"').to_string()
                }
            ))
        } else {
            Err(ParsingErrors {
                record: line.clone(),
                errors: errors
            })
        }
    } else {
        errors.push(ELBLogEntryErrors::MalformedEntry);
        Err(ParsingErrors {
            record: line.clone(),
            errors: errors
        })
    }
}

fn parse_property<T>(raw_prop: &str, prop_name: &'static str, errors: &mut Vec<ELBLogEntryErrors>) -> Option<T>
    where T: FromStr,
    T::Err: Error + 'static,
{
    match raw_prop.parse::<T>() {
        Ok(parsed) => Some(parsed),

        Err(e) => {
            errors.push(
                ELBLogEntryErrors::ParsingError {
                    property: prop_name,
                    description: e.description().to_string(),
                }
            );
            None
        }
    }
}

//TODO test the error case.
//TODO unite the nomenclature under entry (vs record)

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

    const TEST_LINE: &'static str = "2015-08-15T23:43:05.302180Z elb-name 172.16.1.6:54814 \
    172.16.1.5:9000 0.000039 0.145507 0.00003 200 200 0 7582 \
    \"GET http://some.domain.com:80/path0/path1?param0=p0&param1=p1 HTTP/1.1\"\
    ";

    #[test]
	fn parse_line_returns_a_malformed_record_error_for_records_short_on_values() {
        let short_record = "2015-08-15T23:43:05.302180Z elb-name 172.16.1.6:54814 \
        172.16.1.5:9000 0.000039 200 200 0 7582 \
        \"GET http://some.domain.com:80/path0/path1?param0=p0&param1=p1 HTTP/1.1\" \
        ";

        let malformed_error = parse_line(&short_record.to_string()).unwrap_err().errors.pop();

		assert_eq!(malformed_error, Some(ELBLogEntryErrors::MalformedEntry))
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_http_version() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_http_version, "HTTP/1.1")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_url() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_url, "http://some.domain.com:80/path0/path1?param0=p0&param1=p1")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_method() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_method, "GET")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_sent_bytes() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.sent_bytes, 7582)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_received_bytes() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.received_bytes, 0)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_backend_status_code() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.backend_status_code, 200)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_elb_status_code() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.elb_status_code, 200)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_response_processing_time() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.response_processing_time, 0.00003)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_backend_processing_time() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.backend_processing_time, 0.145507)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_request_processing_time() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.request_processing_time, 0.000039)
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_backend_address() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.backend_address, "172.16.1.5:9000".parse().unwrap())
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_client_address() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.client_address, "172.16.1.6:54814".parse().unwrap())
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_timestamp() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(format!("{:?}", elb_log_entry.timestamp), "2015-08-15T23:43:05.302180Z")
	}

    #[test]
	fn parse_line_returns_a_log_entry_with_the_elb_name() {
        let elb_log_entry = parse_line(&TEST_LINE.to_string()).unwrap();

		assert_eq!(elb_log_entry.elb_name, "elb-name")
	}
}
<|endoftext|>


<commit_before>use std::thread;
use std::cell::UnsafeCell;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver, RecvError};
use basetypes::*;
use chess_move::*;
use tt::*;
use position::Position;

pub enum Command {
    Search {
        search_id: usize,
        position: Position,
        depth: u8,
        lower_bound: Value,
        upper_bound: Value,
    },
    Stop,
    Exit,
}

pub enum Report {
    Progress {
        search_id: usize,
        searched_nodes: NodeCount,
        depth: u8,
    },
    Done {
        search_id: usize,
        searched_nodes: NodeCount,
        value: Option<Value>,
    },
}

pub fn run_deepening(tt: Arc<TranspositionTable>,
                     commands: Receiver<Command>,
                     reports: Sender<Report>) {
    let (slave_commands_tx, slave_commands_rx) = channel();
    let (slave_reports_tx, slave_reports_rx) = channel();
    let slave = thread::spawn(move || {
        run(tt, slave_commands_rx, slave_reports_tx);
    });
    let mut pending_command = None;
    loop {
        // If there is a pending command, we take it, otherwise we
        // block and wait to receive a new one.
        let command = match pending_command.take() {
            Some(cmd) => cmd,
            None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
        };

        match command {
            Command::Search { search_id, position, depth, lower_bound, upper_bound } => {
                let mut searched_nodes_final = 0;
                let mut value_final = None;
                'depthloop: for n in 1..(depth + 1) {
                    slave_commands_tx.send(Command::Search {
                                         search_id: n as usize,
                                         position: position.clone(),
                                         depth: n,
                                         lower_bound: lower_bound,
                                         upper_bound: upper_bound,
                                     })
                                     .unwrap();
                    loop {
                        match slave_reports_rx.recv().unwrap() {
                            Report::Progress { searched_nodes, .. } => {
                                reports.send(Report::Progress {
                                           search_id: search_id,
                                           searched_nodes: searched_nodes_final + searched_nodes,
                                           depth: n,
                                       })
                                       .ok();
                                if pending_command.is_none() {
                                    pending_command = match commands.try_recv() {
                                        Ok(cmd) => {
                                            slave_commands_tx.send(Command::Stop).unwrap();
                                            Some(cmd)
                                        }
                                        _ => None,
                                    }
                                }
                            }
                            Report::Done { searched_nodes, value, .. } => {
                                searched_nodes_final += searched_nodes;
                                if n == depth {
                                    value_final = value;
                                }
                                if pending_command.is_some() {
                                    break 'depthloop;
                                }
                                break;
                            }
                        }
                    }
                }
                reports.send(Report::Done {
                           search_id: search_id,
                           searched_nodes: searched_nodes_final,
                           value: value_final,
                       })
                       .ok();
            }
            Command::Stop => {
                slave_commands_tx.send(Command::Stop).unwrap();
                continue;
            }
            Command::Exit => {
                slave_commands_tx.send(Command::Exit).unwrap();
                break;
            }
        }
    }
    slave.join().unwrap();
}

pub fn run(tt: Arc<TranspositionTable>, commands: Receiver<Command>, reports: Sender<Report>) {
    thread_local!(
        static MOVE_STACK: UnsafeCell<MoveStack> = UnsafeCell::new(MoveStack::new())
    );
    MOVE_STACK.with(|s| {
        let mut move_stack = unsafe { &mut *s.get() };
        let mut pending_command = None;
        loop {
            // If there is a pending command, we take it, otherwise we
            // block and wait to receive a new one.
            let command = match pending_command.take() {
                Some(cmd) => cmd,
                None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
            };

            match command {
                Command::Search { search_id, position, depth, lower_bound, upper_bound } => {
                    let mut report = |n| {
                        reports.send(Report::Progress {
                                   search_id: search_id,
                                   searched_nodes: n,
                                   depth: depth,
                               })
                               .ok();
                        match commands.try_recv() {
                            Ok(x) => {
                                pending_command = Some(x);
                                Err(TerminatedSearch)
                            }
                            _ => Ok(()),
                        }
                    };
                    let mut search = Search::new(&tt, position, move_stack, &mut report);
                    let value = search.run(lower_bound, upper_bound, depth).ok();
                    reports.send(Report::Progress {
                               search_id: search_id,
                               searched_nodes: search.node_count(),
                               depth: depth,
                           })
                           .ok();
                    reports.send(Report::Done {
                               search_id: search_id,
                               searched_nodes: search.node_count(),
                               value: value,
                           })
                           .ok();
                    search.reset();
                }
                Command::Stop => continue,
                Command::Exit => break,
            }
        }
    })
}

/// Represents a terminated search condition.
pub struct TerminatedSearch;

enum NodePhase {
    Pristine,
    TriedHashMove,
    GeneratedMoves,
}

struct NodeState {
    phase: NodePhase,
    entry: EntryData,
}

const NODE_COUNT_REPORT_INTERVAL: NodeCount = 10000;

pub struct Search<'a> {
    tt: &'a TranspositionTable,
    position: Position,
    moves: &'a mut MoveStack,
    moves_starting_ply: usize,
    state_stack: Vec<NodeState>,
    reported_nodes: NodeCount,
    unreported_nodes: NodeCount,
    report_function: &'a mut FnMut(NodeCount) -> Result<(), TerminatedSearch>,
}

impl<'a> Search<'a> {
    /// Creates a new instance.
    pub fn new(tt: &'a TranspositionTable,
               position: Position,
               moves: &'a mut MoveStack,
               report_function: &'a mut FnMut(NodeCount) -> Result<(), TerminatedSearch>)
               -> Search<'a> {
        let moves_starting_ply = moves.ply();
        Search {
            tt: tt,
            position: position,
            moves: moves,
            moves_starting_ply: moves_starting_ply,
            state_stack: Vec::with_capacity(32),
            reported_nodes: 0,
            unreported_nodes: 0,
            report_function: report_function,
        }
    }

    /// Performs a principal variation search and returns a result.
    ///
    /// **Important note**: This method may leave un-restored move
    /// lists in the move stack. Call `reset` if you want the move
    /// stack to be restored to the state it had when the search
    /// instance was created.
    pub fn run(&mut self,
               mut alpha: Value, // lower bound
               beta: Value, // upper bound
               depth: u8)
               -> Result<Value, TerminatedSearch> {
        assert!(alpha < beta);

        let entry = self.node_begin();

        // Check if the TT entry gives the result.
        if entry.depth() >= depth {
            let value = entry.value();
            let bound = entry.bound();
            if (value >= beta && bound == BOUND_LOWER) ||
               (value <= alpha && bound == BOUND_UPPER) || (bound == BOUND_EXACT) {
                self.node_end();
                return Ok(value);
            }
        }

        // Initial guests for the final result.
        let mut bound = BOUND_UPPER;
        let mut best_move = Move::invalid();

        if depth == 0 {
            // On leaf nodes, do quiescence search.
            let (value, nodes) = self.position
                                     .evaluate_quiescence(alpha, beta, Some(entry.eval_value()));
            try!(self.report_progress(nodes));

            // See how good this position is.
            if value >= beta {
                alpha = beta;
                bound = BOUND_LOWER;
            } else if value > alpha {
                alpha = value;
                bound = BOUND_EXACT;
            }

        } else {
            // On non-leaf nodes, try moves.
            let mut no_moves_yet = true;
            while let Some(m) = self.do_move() {
                try!(self.report_progress(1));

                // Make a recursive call.
                let value = if no_moves_yet {
                    // The first move we analyze with a fully open window
                    // (alpha, beta). If this happens to be a good move,
                    // it will probably raise `alpha`.
                    no_moves_yet = false;
                    -try!(self.run(-beta, -alpha, depth - 1))
                } else {
                    // For the next moves we first try to prove that they
                    // are not better than our current best move. For this
                    // purpose we analyze them with a null window (alpha,
                    // alpha + 1). This is faster than a full window
                    // search. Only if we are certain that the move is
                    // better than our current best move, we do a
                    // full-window search.
                    match -try!(self.run(-alpha - 1, -alpha, depth - 1)) {
                        x if x <= alpha => x,
                        _ => -try!(self.run(-beta, -alpha, depth - 1)),
                    }
                };
                self.undo_move();

                // See how good this move was.
                if value >= beta {
                    // This move is so good, that the opponent will
                    // probably not allow this line of play to
                    // happen. Therefore we should not lose any more time
                    // on this position.
                    alpha = beta;
                    bound = BOUND_LOWER;
                    best_move = m;
                    break;
                }
                if value > alpha {
                    // We found a new best move.
                    alpha = value;
                    bound = BOUND_EXACT;
                    best_move = m;
                }
            }

            // Check if we are in a final position (no legal moves).
            if no_moves_yet {
                alpha = self.position.evaluate_final();
                bound = BOUND_EXACT;
            }
        }

        self.store(alpha, bound, depth, best_move);
        self.node_end();
        Ok(alpha)
    }

    /// Returns the number of searched positions.
    #[inline(always)]
    pub fn node_count(&self) -> NodeCount {
        self.reported_nodes + self.unreported_nodes
    }

    /// Resets the instance to the state it had when it was created.
    #[inline]
    pub fn reset(&mut self) {
        while self.moves.ply() > self.moves_starting_ply {
            self.moves.restore();
        }
        self.state_stack.clear();
        self.reported_nodes = 0;
        self.unreported_nodes = 0;
    }

    // Declares that we are starting to process a new node.
    //
    // Each recursive call to `run` begins with a call to
    // `node_begin`. The returned value is a TT entry telling
    // everything we know about the current position.
    #[inline]
    fn node_begin(&mut self) -> EntryData {
        // Consult the transposition table.
        let entry = if let Some(e) = self.tt.probe(self.position.hash()) {
            e
        } else {
            EntryData::new(0, BOUND_NONE, 0, 0, self.position.evaluate_static())
        };
        self.state_stack.push(NodeState {
            phase: NodePhase::Pristine,
            entry: entry,
        });
        entry
    }

    // Declares that we are done processing the current node.
    //
    // Each recursive call to `run` ends with a call to `node_end`.
    #[inline]
    fn node_end(&mut self) {
        if let NodePhase::Pristine = self.state_stack.last().unwrap().phase {
            // For pristine nodes we have not saved the move list
            // yet, so we should not restore it.
        } else {
            self.moves.restore();
        }
        self.state_stack.pop();
    }

    // Plays the next legal move in the current position and returns
    // it.
    //
    // Each call to `do_move` for the same position will play and
    // return a different move. When all legal moves has been played,
    // `None` will be returned. `do_move` will do whatever it can to
    // play the best moves first, and the worst last. It will also try
    // to be efficient, for example it will generate the list of all
    // pseudo-legal moves at the last possible moment.
    #[inline]
    fn do_move(&mut self) -> Option<Move> {
        let state = self.state_stack.last_mut().unwrap();

        if let NodePhase::Pristine = state.phase {
            // We save the move list at the last possible moment,
            // because most of the nodes are leafs.
            self.moves.save();

            // We always try the hash move first.
            state.phase = NodePhase::TriedHashMove;
            if state.entry.move16() != 0 {
                if let Some(m) = self.position.try_move_digest(state.entry.move16()) {
                    if self.position.do_move(m) {
                        return Some(m);
                    }
                }
            }
        }

        if let NodePhase::TriedHashMove = state.phase {
            // After the hash move, we generate all pseudo-legal
            // moves. But we should not forget to remove the already
            // tried hash move from the list.
            self.position.generate_moves(self.moves);
            if state.entry.move16() != 0 {
                self.moves.remove_move(state.entry.move16());
            }
            state.phase = NodePhase::GeneratedMoves;
        }

        // For the last, we spit the generated moves out.
        while let Some(m) = self.moves.remove_best_move() {
            if self.position.do_move(m) {
                return Some(m);
            }
        }
        None
    }

    // Takes the last played move back.
    #[inline]
    fn undo_move(&mut self) {
        self.position.undo_move();
    }

    // Stores updated node information in the transposition table.
    #[inline]
    fn store(&mut self, value: Value, bound: BoundType, depth: u8, best_move: Move) {
        let entry = &self.state_stack.last().unwrap().entry;
        let move16 = match best_move.digest() {
            0 => entry.move16(),
            x => x,
        };
        self.tt.store(self.position.hash(),
                      EntryData::new(value, bound, depth, move16, entry.eval_value()));
    }

    // Reports search progress.
    //
    // From time to time, we should report how many nodes had been
    // searched since the beginning of the search. This also gives an
    // opportunity for the search to be terminated.
    #[inline]
    fn report_progress(&mut self, new_nodes: NodeCount) -> Result<(), TerminatedSearch> {
        self.unreported_nodes += new_nodes;
        if self.unreported_nodes > NODE_COUNT_REPORT_INTERVAL {
            self.reported_nodes += self.unreported_nodes;
            self.unreported_nodes = 0;
            try!((*self.report_function)(self.reported_nodes));
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::Search;
    use chess_move::*;
    use tt::*;
    use position::Position;

    #[test]
    fn test_search() {
        let p = Position::from_fen("8/8/8/8/3q3k/7n/6PP/2Q2R1K b - - 0 1").ok().unwrap();
        let tt = TranspositionTable::new();
        let mut moves = MoveStack::new();
        let mut report = |_| Ok(());
        let mut search = Search::new(&tt, p, &mut moves, &mut report);
        let value = search.run(-30000, 30000, 2)
                          .ok()
                          .unwrap();
        assert!(value < -300);
        search.reset();
        let value = search.run(-30000, 30000, 4)
                          .ok()
                          .unwrap();
        assert!(value >= 20000);
    }
}
<commit_msg>beautification<commit_after>use std::thread;
use std::cell::UnsafeCell;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver, RecvError};
use basetypes::*;
use chess_move::*;
use tt::*;
use position::Position;

pub enum Command {
    Search {
        search_id: usize,
        position: Position,
        depth: u8,
        lower_bound: Value,
        upper_bound: Value,
    },
    Stop,
    Exit,
}

pub enum Report {
    Progress {
        search_id: usize,
        searched_nodes: NodeCount,
        depth: u8,
    },
    Done {
        search_id: usize,
        searched_nodes: NodeCount,
        value: Option<Value>,
    },
}

pub fn run_deepening(tt: Arc<TranspositionTable>,
                     commands: Receiver<Command>,
                     reports: Sender<Report>) {
    let (slave_commands_tx, slave_commands_rx) = channel();
    let (slave_reports_tx, slave_reports_rx) = channel();
    let slave = thread::spawn(move || {
        run(tt, slave_commands_rx, slave_reports_tx);
    });
    let mut pending_command = None;
    loop {
        // If there is a pending command, we take it, otherwise we
        // block and wait to receive a new one.
        let command = match pending_command.take() {
            Some(cmd) => cmd,
            None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
        };

        match command {
            Command::Search { search_id, position, depth, lower_bound, upper_bound } => {
                let mut searched_nodes_final = 0;
                let mut value_final = None;
                'depthloop: for n in 1..(depth + 1) {
                    slave_commands_tx.send(Command::Search {
                                         search_id: n as usize,
                                         position: position.clone(),
                                         depth: n,
                                         lower_bound: lower_bound,
                                         upper_bound: upper_bound,
                                     })
                                     .unwrap();
                    loop {
                        match slave_reports_rx.recv().unwrap() {
                            Report::Progress { searched_nodes, .. } => {
                                reports.send(Report::Progress {
                                           search_id: search_id,
                                           searched_nodes: searched_nodes_final + searched_nodes,
                                           depth: n,
                                       })
                                       .ok();
                                if pending_command.is_none() {
                                    pending_command = match commands.try_recv() {
                                        Ok(cmd) => {
                                            slave_commands_tx.send(Command::Stop).unwrap();
                                            Some(cmd)
                                        }
                                        _ => None,
                                    }
                                }
                            }
                            Report::Done { searched_nodes, value, .. } => {
                                searched_nodes_final += searched_nodes;
                                if n == depth {
                                    value_final = value;
                                }
                                if pending_command.is_some() {
                                    break 'depthloop;
                                }
                                break;
                            }
                        }
                    }
                }
                reports.send(Report::Done {
                           search_id: search_id,
                           searched_nodes: searched_nodes_final,
                           value: value_final,
                       })
                       .ok();
            }
            Command::Stop => {
                slave_commands_tx.send(Command::Stop).unwrap();
                continue;
            }
            Command::Exit => {
                slave_commands_tx.send(Command::Exit).unwrap();
                break;
            }
        }
    }
    slave.join().unwrap();
}

pub fn run(tt: Arc<TranspositionTable>, commands: Receiver<Command>, reports: Sender<Report>) {
    thread_local!(
        static MOVE_STACK: UnsafeCell<MoveStack> = UnsafeCell::new(MoveStack::new())
    );
    MOVE_STACK.with(|s| {
        let mut move_stack = unsafe { &mut *s.get() };
        let mut pending_command = None;
        loop {
            // If there is a pending command, we take it, otherwise we
            // block and wait to receive a new one.
            let command = match pending_command.take() {
                Some(cmd) => cmd,
                None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
            };

            match command {
                Command::Search { search_id, position, depth, lower_bound, upper_bound } => {
                    let mut report = |n| {
                        reports.send(Report::Progress {
                                   search_id: search_id,
                                   searched_nodes: n,
                                   depth: depth,
                               })
                               .ok();
                        match commands.try_recv() {
                            Ok(x) => {
                                pending_command = Some(x);
                                Err(TerminatedSearch)
                            }
                            _ => Ok(()),
                        }
                    };
                    let mut search = Search::new(position, &tt, move_stack, &mut report);
                    let value = search.run(lower_bound, upper_bound, depth).ok();
                    reports.send(Report::Progress {
                               search_id: search_id,
                               searched_nodes: search.node_count(),
                               depth: depth,
                           })
                           .ok();
                    reports.send(Report::Done {
                               search_id: search_id,
                               searched_nodes: search.node_count(),
                               value: value,
                           })
                           .ok();
                    search.reset();
                }
                Command::Stop => continue,
                Command::Exit => break,
            }
        }
    })
}

/// Represents a terminated search condition.
pub struct TerminatedSearch;

enum NodePhase {
    Pristine,
    TriedHashMove,
    GeneratedMoves,
}

struct NodeState {
    phase: NodePhase,
    entry: EntryData,
}

const NODE_COUNT_REPORT_INTERVAL: NodeCount = 10000;

pub struct Search<'a> {
    tt: &'a TranspositionTable,
    position: Position,
    moves: &'a mut MoveStack,
    moves_starting_ply: usize,
    state_stack: Vec<NodeState>,
    reported_nodes: NodeCount,
    unreported_nodes: NodeCount,
    report_function: &'a mut FnMut(NodeCount) -> Result<(), TerminatedSearch>,
}

impl<'a> Search<'a> {
    /// Creates a new instance.
    pub fn new(p: Position,
               tt: &'a TranspositionTable,
               moves: &'a mut MoveStack,
               report_function: &'a mut FnMut(NodeCount) -> Result<(), TerminatedSearch>)
               -> Search<'a> {
        let moves_starting_ply = moves.ply();
        Search {
            tt: tt,
            position: p,
            moves: moves,
            moves_starting_ply: moves_starting_ply,
            state_stack: Vec::with_capacity(32),
            reported_nodes: 0,
            unreported_nodes: 0,
            report_function: report_function,
        }
    }

    /// Performs a principal variation search and returns a result.
    ///
    /// **Important note**: This method may leave un-restored move
    /// lists in the move stack. Call `reset` if you want the move
    /// stack to be restored to the state it had when the search
    /// instance was created.
    pub fn run(&mut self,
               mut alpha: Value, // lower bound
               beta: Value, // upper bound
               depth: u8)
               -> Result<Value, TerminatedSearch> {
        assert!(alpha < beta);

        let entry = self.node_begin();

        // Check if the TT entry gives the result.
        if entry.depth() >= depth {
            let value = entry.value();
            let bound = entry.bound();
            if (value >= beta && bound == BOUND_LOWER) ||
               (value <= alpha && bound == BOUND_UPPER) || (bound == BOUND_EXACT) {
                self.node_end();
                return Ok(value);
            }
        }

        // Initial guests for the final result.
        let mut bound = BOUND_UPPER;
        let mut best_move = Move::invalid();

        if depth == 0 {
            // On leaf nodes, do quiescence search.
            let (value, nodes) = self.position
                                     .evaluate_quiescence(alpha, beta, Some(entry.eval_value()));
            try!(self.report_progress(nodes));

            // See how good this position is.
            if value >= beta {
                alpha = beta;
                bound = BOUND_LOWER;
            } else if value > alpha {
                alpha = value;
                bound = BOUND_EXACT;
            }

        } else {
            // On non-leaf nodes, try moves.
            let mut no_moves_yet = true;
            while let Some(m) = self.do_move() {
                try!(self.report_progress(1));

                // Make a recursive call.
                let value = if no_moves_yet {
                    // The first move we analyze with a fully open window
                    // (alpha, beta). If this happens to be a good move,
                    // it will probably raise `alpha`.
                    no_moves_yet = false;
                    -try!(self.run(-beta, -alpha, depth - 1))
                } else {
                    // For the next moves we first try to prove that they
                    // are not better than our current best move. For this
                    // purpose we analyze them with a null window (alpha,
                    // alpha + 1). This is faster than a full window
                    // search. Only if we are certain that the move is
                    // better than our current best move, we do a
                    // full-window search.
                    match -try!(self.run(-alpha - 1, -alpha, depth - 1)) {
                        x if x <= alpha => x,
                        _ => -try!(self.run(-beta, -alpha, depth - 1)),
                    }
                };
                self.undo_move();

                // See how good this move was.
                if value >= beta {
                    // This move is so good, that the opponent will
                    // probably not allow this line of play to
                    // happen. Therefore we should not lose any more time
                    // on this position.
                    alpha = beta;
                    bound = BOUND_LOWER;
                    best_move = m;
                    break;
                }
                if value > alpha {
                    // We found a new best move.
                    alpha = value;
                    bound = BOUND_EXACT;
                    best_move = m;
                }
            }

            // Check if we are in a final position (no legal moves).
            if no_moves_yet {
                alpha = self.position.evaluate_final();
                bound = BOUND_EXACT;
            }
        }

        self.store(alpha, bound, depth, best_move);
        self.node_end();
        Ok(alpha)
    }

    /// Returns the number of searched positions.
    #[inline(always)]
    pub fn node_count(&self) -> NodeCount {
        self.reported_nodes + self.unreported_nodes
    }

    /// Resets the instance to the state it had when it was created.
    #[inline]
    pub fn reset(&mut self) {
        while self.moves.ply() > self.moves_starting_ply {
            self.moves.restore();
        }
        self.state_stack.clear();
        self.reported_nodes = 0;
        self.unreported_nodes = 0;
    }

    // Declares that we are starting to process a new node.
    //
    // Each recursive call to `run` begins with a call to
    // `node_begin`. The returned value is a TT entry telling
    // everything we know about the current position.
    #[inline]
    fn node_begin(&mut self) -> EntryData {
        // Consult the transposition table.
        let entry = if let Some(e) = self.tt.probe(self.position.hash()) {
            e
        } else {
            EntryData::new(0, BOUND_NONE, 0, 0, self.position.evaluate_static())
        };
        self.state_stack.push(NodeState {
            phase: NodePhase::Pristine,
            entry: entry,
        });
        entry
    }

    // Declares that we are done processing the current node.
    //
    // Each recursive call to `run` ends with a call to `node_end`.
    #[inline]
    fn node_end(&mut self) {
        if let NodePhase::Pristine = self.state_stack.last().unwrap().phase {
            // For pristine nodes we have not saved the move list
            // yet, so we should not restore it.
        } else {
            self.moves.restore();
        }
        self.state_stack.pop();
    }

    // Plays the next legal move in the current position and returns
    // it.
    //
    // Each call to `do_move` for the same position will play and
    // return a different move. When all legal moves has been played,
    // `None` will be returned. `do_move` will do whatever it can to
    // play the best moves first, and the worst last. It will also try
    // to be efficient, for example it will generate the list of all
    // pseudo-legal moves at the last possible moment.
    #[inline]
    fn do_move(&mut self) -> Option<Move> {
        let state = self.state_stack.last_mut().unwrap();

        if let NodePhase::Pristine = state.phase {
            // We save the move list at the last possible moment,
            // because most of the nodes are leafs.
            self.moves.save();

            // We always try the hash move first.
            state.phase = NodePhase::TriedHashMove;
            if state.entry.move16() != 0 {
                if let Some(m) = self.position.try_move_digest(state.entry.move16()) {
                    if self.position.do_move(m) {
                        return Some(m);
                    }
                }
            }
        }

        if let NodePhase::TriedHashMove = state.phase {
            // After the hash move, we generate all pseudo-legal
            // moves. But we should not forget to remove the already
            // tried hash move from the list.
            self.position.generate_moves(self.moves);
            if state.entry.move16() != 0 {
                self.moves.remove_move(state.entry.move16());
            }
            state.phase = NodePhase::GeneratedMoves;
        }

        // For the last, we spit the generated moves out.
        while let Some(m) = self.moves.remove_best_move() {
            if self.position.do_move(m) {
                return Some(m);
            }
        }
        None
    }

    // Takes the last played move back.
    #[inline]
    fn undo_move(&mut self) {
        self.position.undo_move();
    }

    // Stores updated node information in the transposition table.
    #[inline]
    fn store(&mut self, value: Value, bound: BoundType, depth: u8, best_move: Move) {
        let entry = &self.state_stack.last().unwrap().entry;
        let move16 = match best_move.digest() {
            0 => entry.move16(),
            x => x,
        };
        self.tt.store(self.position.hash(),
                      EntryData::new(value, bound, depth, move16, entry.eval_value()));
    }

    // Reports search progress.
    //
    // From time to time, we should report how many nodes had been
    // searched since the beginning of the search. This also gives an
    // opportunity for the search to be terminated.
    #[inline]
    fn report_progress(&mut self, new_nodes: NodeCount) -> Result<(), TerminatedSearch> {
        self.unreported_nodes += new_nodes;
        if self.unreported_nodes > NODE_COUNT_REPORT_INTERVAL {
            self.reported_nodes += self.unreported_nodes;
            self.unreported_nodes = 0;
            try!((*self.report_function)(self.reported_nodes));
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::Search;
    use chess_move::*;
    use tt::*;
    use position::Position;

    #[test]
    fn test_search() {
        let p = Position::from_fen("8/8/8/8/3q3k/7n/6PP/2Q2R1K b - - 0 1").ok().unwrap();
        let tt = TranspositionTable::new();
        let mut moves = MoveStack::new();
        let mut report = |_| Ok(());
        let mut search = Search::new(p, &tt, &mut moves, &mut report);
        let value = search.run(-30000, 30000, 2)
                          .ok()
                          .unwrap();
        assert!(value < -300);
        search.reset();
        let value = search.run(-30000, 30000, 4)
                          .ok()
                          .unwrap();
        assert!(value >= 20000);
    }
}
<|endoftext|>


<commit_before>use std::fmt;

use crate::error::RenderingError;

/// An enumeration of errors that can occur during filter primitive rendering.
#[derive(Debug, Clone, PartialEq)]
pub enum FilterError {
    /// The units on the filter bounds are invalid
    InvalidUnits,
    /// The filter was passed invalid input (the `in` attribute).
    InvalidInput,
    /// The filter input surface has an unsuccessful status.
    BadInputSurfaceStatus(cairo::Status),
    /// A Cairo error.
    ///
    /// This means that either a failed intermediate surface creation or bad intermediate surface
    /// status.
    CairoError(cairo::Status),
    /// Error from the rendering backend.
    Rendering(RenderingError),
    /// A lighting filter has none or multiple light sources.
    InvalidLightSourceCount,
    /// A lighting filter input surface is too small.
    LightingInputTooSmall,
    /// Child node was in error.
    ChildNodeInError,
}

impl fmt::Display for FilterError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            FilterError::InvalidUnits => write!(
                f,
                "unit identifiers are not allowed with primitiveUnits set to objectBoundingBox"
            ),
            FilterError::InvalidInput => write!(f, "invalid value of the `in` attribute"),
            FilterError::BadInputSurfaceStatus(ref status) => {
                write!(f, "invalid status of the input surface: {}", status)
            }
            FilterError::CairoError(ref status) => write!(f, "Cairo error: {}", status),
            FilterError::Rendering(ref e) => write!(f, "Rendering error: {}", e),
            FilterError::InvalidLightSourceCount => write!(f, "invalid light source count"),
            FilterError::LightingInputTooSmall => write!(
                f,
                "lighting filter input surface is too small (less than 2×2 pixels)"
            ),
            FilterError::ChildNodeInError => write!(f, "child node was in error"),
        }
    }
}

impl From<cairo::Status> for FilterError {
    #[inline]
    fn from(x: cairo::Status) -> Self {
        FilterError::CairoError(x)
    }
}

impl From<RenderingError> for FilterError {
    #[inline]
    fn from(e: RenderingError) -> Self {
        FilterError::Rendering(e)
    }
}
<commit_msg>FilterError does not need to be PartialEq<commit_after>use std::fmt;

use crate::error::RenderingError;

/// An enumeration of errors that can occur during filter primitive rendering.
#[derive(Debug, Clone)]
pub enum FilterError {
    /// The units on the filter bounds are invalid
    InvalidUnits,
    /// The filter was passed invalid input (the `in` attribute).
    InvalidInput,
    /// The filter input surface has an unsuccessful status.
    BadInputSurfaceStatus(cairo::Status),
    /// A Cairo error.
    ///
    /// This means that either a failed intermediate surface creation or bad intermediate surface
    /// status.
    CairoError(cairo::Status),
    /// Error from the rendering backend.
    Rendering(RenderingError),
    /// A lighting filter has none or multiple light sources.
    InvalidLightSourceCount,
    /// A lighting filter input surface is too small.
    LightingInputTooSmall,
    /// Child node was in error.
    ChildNodeInError,
}

impl fmt::Display for FilterError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            FilterError::InvalidUnits => write!(
                f,
                "unit identifiers are not allowed with primitiveUnits set to objectBoundingBox"
            ),
            FilterError::InvalidInput => write!(f, "invalid value of the `in` attribute"),
            FilterError::BadInputSurfaceStatus(ref status) => {
                write!(f, "invalid status of the input surface: {}", status)
            }
            FilterError::CairoError(ref status) => write!(f, "Cairo error: {}", status),
            FilterError::Rendering(ref e) => write!(f, "Rendering error: {}", e),
            FilterError::InvalidLightSourceCount => write!(f, "invalid light source count"),
            FilterError::LightingInputTooSmall => write!(
                f,
                "lighting filter input surface is too small (less than 2×2 pixels)"
            ),
            FilterError::ChildNodeInError => write!(f, "child node was in error"),
        }
    }
}

impl From<cairo::Status> for FilterError {
    #[inline]
    fn from(x: cairo::Status) -> Self {
        FilterError::CairoError(x)
    }
}

impl From<RenderingError> for FilterError {
    #[inline]
    fn from(e: RenderingError) -> Self {
        FilterError::Rendering(e)
    }
}
<|endoftext|>


<commit_before>//! I is the integral image of an image F if I(x, y) is the
//! sum of F(x', y') for x' <= x, y' <= y. i.e. each pixel
//! in the integral image contains the sum of the pixel intensities
//! of all input pixels that are above it and to its left.
//! The integral image has the helpful property that it lets us
//! compute the sum of pixel intensities from any rectangular region
//! in the input image in constant time.
//! Specifically, given a rectangle in F with clockwise corners
//! A, B, C, D, with A at the upper left, the total pixel intensity
//! of this rectangle is I(C) - I(B) - I(D) + I(A).

extern crate image;

use image::{
    Luma,
    GenericImage,
    ImageBuffer
};

/// Computes the integral image of an 8bpp grayscale image.
// TODO: Support more formats.
// TODO: This is extremely slow. Fix that!
pub fn integral_image<I: GenericImage<Pixel=Luma<u8>> + 'static>(image: &I)
    -> ImageBuffer<Luma<u32>, Vec<u32>> {
    padded_integral_image(image, 0, 0)
}

/// Computes the integral image of the result of padding image
/// with its boundary pixels for x_padding columns on either
/// side and y_padding rows at its top and bottom.
/// Returned image has width image.width() + 2 * x_padding
/// and height image.height() + 2 * y_padding.
pub fn padded_integral_image<I: GenericImage<Pixel=Luma<u8>> + 'static>(
    image: &I,
    x_padding: u32,
    y_padding: u32)
    -> ImageBuffer<Luma<u32>, Vec<u32>> {

    let (in_width, in_height) = image.dimensions();
    let out_width = in_width + 2 * x_padding;
    let out_height = in_height + 2 * y_padding;

    let mut out: ImageBuffer<Luma<u32>, Vec<u32>>
        = ImageBuffer::new(out_width, out_height);

    for y in 0..out_height {
        for x in 0..out_width {

            let y_in: u32;
            if y < y_padding {
                y_in = 0;
            }
            else if y >= in_height + y_padding {
                y_in = in_height - 1;
            }
            else {
                y_in = y - y_padding;
            }
            let x_in: u32;
            if x < x_padding {
                x_in = 0;
            }
            else if x >= in_width + x_padding {
                x_in = in_width - 1;
            }
            else {
                x_in = x - x_padding;
            }

            let p = image.get_pixel(x_in, y_in);
            out.put_pixel(x, y, Luma([p[0] as u32]));
        }
    }

    for x in 1..out_width {
        (*out.get_pixel_mut(x, 0))[0] += out.get_pixel(x - 1, 0)[0];
    }

    for y in 1..out_height {
        (*out.get_pixel_mut(0, y))[0] += out.get_pixel(0, y - 1)[0];

        for x in 1..out_width {
            (*out.get_pixel_mut(x, y))[0] += out.get_pixel(x, y - 1)[0];
            (*out.get_pixel_mut(x, y))[0] += out.get_pixel(x - 1, y)[0];
            (*out.get_pixel_mut(x, y))[0] -= out.get_pixel(x - 1, y - 1)[0];
        }
    }

    out
}

/// Computes the running sum of one row of image, padded
/// at the beginning and end. The padding is by continuity.
/// Takes a reference to buffer so that this can be reused
/// for all rows in an image.
// TODO: faster, more formats
pub fn row_running_sum<I: GenericImage<Pixel=Luma<u8>> + 'static>(
    image: &I, row: u32, buffer: &mut [u32], padding: u32) {

    let width = image.width();
    assert!(buffer.len() >= (width + 2 * padding) as usize,
        format!("Buffer length {} is less than 2 * {} + {}",
            buffer.len(), width, padding));

    for x in 0..padding {
        buffer[x as usize] = image.get_pixel(0, row)[0] as u32;
    }

    for x in 0..width {
        let idx = (x + padding) as usize;
        buffer[idx] = image.get_pixel(x, row)[0] as u32;
    }

    for x in 0..padding {
        let idx = (x + width + padding) as usize;
        buffer[idx] = image.get_pixel(width - 1, row)[0] as u32;
    }

    for x in 1..width + 2 * padding {
        buffer[x as usize] += buffer[(x - 1) as usize] as u32;
    }
}

#[cfg(test)]
mod test {

    use super::{
        integral_image,
        padded_integral_image,
        row_running_sum
    };
    use utils::{
        gray_bench_image
    };
    use image::{
        GrayImage,
        ImageBuffer,
        Luma
    };
    use test;

    #[test]
    fn test_integral_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 2, vec![
            1u8, 2u8, 3u8,
            4u8, 5u8, 6u8]).unwrap();

        let expected: ImageBuffer<Luma<u32>, Vec<u32>>
            = ImageBuffer::from_raw(3, 2, vec![
            1u32,  3u32,  6u32,
            5u32, 12u32, 21u32]).unwrap();

        assert_pixels_eq!(integral_image(&image), expected);
    }

    #[bench]
    fn bench_integral_image(b: &mut test::Bencher) {
        let image = gray_bench_image(500, 500);
        b.iter(|| {
            let integral = integral_image(&image);
            test::black_box(integral);
            });
    }

    #[test]
    fn test_padded_integral_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 2, vec![
            1u8, 2u8, 3u8,
            4u8, 5u8, 6u8]).unwrap();

        let expected: ImageBuffer<Luma<u32>, Vec<u32>>
            = ImageBuffer::from_raw(5, 4, vec![
              1u32,  2u32,   4u32,   7u32,  10u32,
              2u32,  4u32,   8u32,  14u32,  20u32,
              6u32, 12u32,  21u32,  33u32,  45u32,
             10u32, 20u32,  34u32,  52u32,  70u32]).unwrap();

        assert_pixels_eq!(padded_integral_image(&image, 1, 1), expected);
    }

    #[test]
    fn test_row_running_sum() {
        let image: GrayImage = ImageBuffer::from_raw(3, 2, vec![
            1u8, 2u8, 3u8,
            4u8, 5u8, 6u8]).unwrap();

        let expected = [1, 2, 4, 7, 10];

        let mut buffer = [0; 5];
        row_running_sum(&image, 0, &mut buffer, 1);

        assert_eq!(buffer, expected);
    }

    #[bench]
    fn bench_row_running_sum(b: &mut test::Bencher) {
        let image = gray_bench_image(1000, 1);
        let mut buffer = [0; 1010];
        b.iter(|| {
            row_running_sum(&image, 0, &mut buffer, 5);
            });
    }
}
<commit_msg>Add column_running_sum (which is suspiciously similar to row_running_sum)<commit_after>//! I is the integral image of an image F if I(x, y) is the
//! sum of F(x', y') for x' <= x, y' <= y. i.e. each pixel
//! in the integral image contains the sum of the pixel intensities
//! of all input pixels that are above it and to its left.
//! The integral image has the helpful property that it lets us
//! compute the sum of pixel intensities from any rectangular region
//! in the input image in constant time.
//! Specifically, given a rectangle in F with clockwise corners
//! A, B, C, D, with A at the upper left, the total pixel intensity
//! of this rectangle is I(C) - I(B) - I(D) + I(A).

extern crate image;

use image::{
    Luma,
    GenericImage,
    ImageBuffer
};

/// Computes the integral image of an 8bpp grayscale image.
// TODO: Support more formats.
// TODO: This is extremely slow. Fix that!
pub fn integral_image<I: GenericImage<Pixel=Luma<u8>> + 'static>(image: &I)
    -> ImageBuffer<Luma<u32>, Vec<u32>> {
    padded_integral_image(image, 0, 0)
}

/// Computes the integral image of the result of padding image
/// with its boundary pixels for x_padding columns on either
/// side and y_padding rows at its top and bottom.
/// Returned image has width image.width() + 2 * x_padding
/// and height image.height() + 2 * y_padding.
pub fn padded_integral_image<I: GenericImage<Pixel=Luma<u8>> + 'static>(
    image: &I,
    x_padding: u32,
    y_padding: u32)
    -> ImageBuffer<Luma<u32>, Vec<u32>> {

    let (in_width, in_height) = image.dimensions();
    let out_width = in_width + 2 * x_padding;
    let out_height = in_height + 2 * y_padding;

    let mut out: ImageBuffer<Luma<u32>, Vec<u32>>
        = ImageBuffer::new(out_width, out_height);

    for y in 0..out_height {
        for x in 0..out_width {

            let y_in: u32;
            if y < y_padding {
                y_in = 0;
            }
            else if y >= in_height + y_padding {
                y_in = in_height - 1;
            }
            else {
                y_in = y - y_padding;
            }
            let x_in: u32;
            if x < x_padding {
                x_in = 0;
            }
            else if x >= in_width + x_padding {
                x_in = in_width - 1;
            }
            else {
                x_in = x - x_padding;
            }

            let p = image.get_pixel(x_in, y_in);
            out.put_pixel(x, y, Luma([p[0] as u32]));
        }
    }

    for x in 1..out_width {
        (*out.get_pixel_mut(x, 0))[0] += out.get_pixel(x - 1, 0)[0];
    }

    for y in 1..out_height {
        (*out.get_pixel_mut(0, y))[0] += out.get_pixel(0, y - 1)[0];

        for x in 1..out_width {
            (*out.get_pixel_mut(x, y))[0] += out.get_pixel(x, y - 1)[0];
            (*out.get_pixel_mut(x, y))[0] += out.get_pixel(x - 1, y)[0];
            (*out.get_pixel_mut(x, y))[0] -= out.get_pixel(x - 1, y - 1)[0];
        }
    }

    out
}

/// Computes the running sum of one row of image, padded
/// at the beginning and end. The padding is by continuity.
/// Takes a reference to buffer so that this can be reused
/// for all rows in an image.
// TODO: faster, more formats
pub fn row_running_sum<I: GenericImage<Pixel=Luma<u8>> + 'static>(
    image: &I, row: u32, buffer: &mut [u32], padding: u32) {

    let width = image.width();
    assert!(buffer.len() >= (width + 2 * padding) as usize,
        format!("Buffer length {} is less than 2 * {} + {}",
            buffer.len(), width, padding));

    for x in 0..padding {
        buffer[x as usize] = image.get_pixel(0, row)[0] as u32;
    }

    for x in 0..width {
        let idx = (x + padding) as usize;
        buffer[idx] = image.get_pixel(x, row)[0] as u32;
    }

    for x in 0..padding {
        let idx = (x + width + padding) as usize;
        buffer[idx] = image.get_pixel(width - 1, row)[0] as u32;
    }

    for x in 1..width + 2 * padding {
        buffer[x as usize] += buffer[(x - 1) as usize] as u32;
    }
}

/// Computes the running sum of one column of image, padded
/// at the top and bottom. The padding is by continuity.
/// Takes a reference to buffer so that this can be reused
/// for all columns in an image.
// TODO: faster, more formats
pub fn column_running_sum<I: GenericImage<Pixel=Luma<u8>> + 'static>(
    image: &I, column: u32, buffer: &mut [u32], padding: u32) {

    let height = image.height();
    assert!(buffer.len() >= (height + 2 * padding) as usize,
        format!("Buffer length {} is less than 2 * {} + {}",
            buffer.len(), height, padding));

    for y in 0..padding {
        buffer[y as usize] = image.get_pixel(column, 0)[0] as u32;
    }

    for y in 0..height {
        let idx = (y + padding) as usize;
        buffer[idx] = image.get_pixel(column, y)[0] as u32;
    }

    for y in 0..padding {
        let idx = (y + height + padding) as usize;
        buffer[idx] = image.get_pixel(column, height - 1)[0] as u32;
    }

    for y in 1..height + 2 * padding {
        buffer[y as usize] += buffer[(y - 1) as usize] as u32;
    }
}

#[cfg(test)]
mod test {

    use super::{
        column_running_sum,
        integral_image,
        padded_integral_image,
        row_running_sum
    };
    use utils::{
        gray_bench_image
    };
    use image::{
        GrayImage,
        ImageBuffer,
        Luma
    };
    use test;

    #[test]
    fn test_integral_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 2, vec![
            1u8, 2u8, 3u8,
            4u8, 5u8, 6u8]).unwrap();

        let expected: ImageBuffer<Luma<u32>, Vec<u32>>
            = ImageBuffer::from_raw(3, 2, vec![
            1u32,  3u32,  6u32,
            5u32, 12u32, 21u32]).unwrap();

        assert_pixels_eq!(integral_image(&image), expected);
    }

    #[bench]
    fn bench_integral_image(b: &mut test::Bencher) {
        let image = gray_bench_image(500, 500);
        b.iter(|| {
            let integral = integral_image(&image);
            test::black_box(integral);
            });
    }

    #[test]
    fn test_padded_integral_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 2, vec![
            1u8, 2u8, 3u8,
            4u8, 5u8, 6u8]).unwrap();

        let expected: ImageBuffer<Luma<u32>, Vec<u32>>
            = ImageBuffer::from_raw(5, 4, vec![
              1u32,  2u32,   4u32,   7u32,  10u32,
              2u32,  4u32,   8u32,  14u32,  20u32,
              6u32, 12u32,  21u32,  33u32,  45u32,
             10u32, 20u32,  34u32,  52u32,  70u32]).unwrap();

        assert_pixels_eq!(padded_integral_image(&image, 1, 1), expected);
    }

    #[test]
    fn test_row_running_sum() {
        let image: GrayImage = ImageBuffer::from_raw(3, 2, vec![
            1u8, 2u8, 3u8,
            4u8, 5u8, 6u8]).unwrap();

        let expected = [1, 2, 4, 7, 10];

        let mut buffer = [0; 5];
        row_running_sum(&image, 0, &mut buffer, 1);

        assert_eq!(buffer, expected);
    }

    #[bench]
    fn bench_row_running_sum(b: &mut test::Bencher) {
        let image = gray_bench_image(1000, 1);
        let mut buffer = [0; 1010];
        b.iter(|| {
            row_running_sum(&image, 0, &mut buffer, 5);
            });
    }

    #[test]
    fn test_column_running_sum() {
        let image: GrayImage = ImageBuffer::from_raw(2, 3, vec![
            1u8, 4u8,
            2u8, 5u8,
            3u8, 6u8]).unwrap();

        let expected = [1, 2, 4, 7, 10];

        let mut buffer = [0; 5];
        column_running_sum(&image, 0, &mut buffer, 1);

        assert_eq!(buffer, expected);
    }

    #[bench]
    fn bench_column_running_sum(b: &mut test::Bencher) {
        let image = gray_bench_image(100, 1000);
        let mut buffer = [0; 1010];
        b.iter(|| {
            column_running_sum(&image, 0, &mut buffer, 5);
            });
    }
}
<|endoftext|>