<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Boden Garman <bpgarman@gmail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

// spell-checker:ignore (ToDO) fpath

extern crate clap;
#[macro_use]
extern crate uucore;

use clap::{App, Arg, ArgMatches};

use std::fs::File;
use std::io::{stdin, BufRead, BufReader, Read};
use std::path::Path;
use std::result::Result as StdResult;
use std::str::from_utf8;

struct Settings {
    show_bytes: bool,
    show_chars: bool,
    show_lines: bool,
    show_words: bool,
    show_max_line_length: bool,
}

impl Settings {
    fn new(matches: &ArgMatches) -> Settings {
        let settings = Settings {
            show_bytes: matches.is_present(OPT_BYTES),
            show_chars: matches.is_present(OPT_CHAR),
            show_lines: matches.is_present(OPT_LINES),
            show_words: matches.is_present(OPT_WORDS),
            show_max_line_length: matches.is_present(OPT_MAX_LINE_LENGTH),
        };

        if settings.show_bytes
            || settings.show_chars
            || settings.show_lines
            || settings.show_words
            || settings.show_max_line_length
        {
            return settings;
        }

        Settings {
            show_bytes: true,
            show_chars: false,
            show_lines: true,
            show_words: true,
            show_max_line_length: false,
        }
    }
}

struct Result {
    title: String,
    bytes: usize,
    chars: usize,
    lines: usize,
    words: usize,
    max_line_length: usize,
}

static ABOUT: &str = "print newline, word, and byte counts for each file.";
static VERSION: &str = env!("CARGO_PKG_VERSION");

static OPT_BYTES: &str = "bytes";
static OPT_CHAR: &str = "chars";
static OPT_LINES: &str = "lines";
static OPT_MAX_LINE_LENGTH: &str = "max-line-length";
static OPT_WORDS: &str = "words";

static ARG_FILES: &str = "files";

fn get_usage() -> String {
    format!(
        "{0} [OPTION]... [FILE]...
 With no FILE, or when FILE is -, read standard input.",
        executable!()
    )
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let usage = get_usage();

    let matches = App::new(executable!())
        .version(VERSION)
        .about(ABOUT)
        .usage(&usage[..])
        .arg(
            Arg::with_name(OPT_BYTES)
                .short("c")
                .long(OPT_BYTES)
                .help("print the byte counts"),
        )
        .arg(
            Arg::with_name(OPT_CHAR)
                .short("m")
                .long(OPT_CHAR)
                .help("print the character counts"),
        )
        .arg(
            Arg::with_name(OPT_LINES)
                .short("l")
                .long(OPT_LINES)
                .help("print the newline counts"),
        )
        .arg(
            Arg::with_name(OPT_MAX_LINE_LENGTH)
                .short("L")
                .long(OPT_MAX_LINE_LENGTH)
                .help("print the length of the longest line"),
        )
        .arg(
            Arg::with_name(OPT_WORDS)
                .short("w")
                .long(OPT_WORDS)
                .help("print the word counts"),
        )
        .arg(Arg::with_name(ARG_FILES).multiple(true).takes_value(true))
        .get_matches_from(args);

    let mut files: Vec<String> = matches
        .values_of(ARG_FILES)
        .map(|v| v.map(ToString::to_string).collect())
        .unwrap_or_default();

    if files.is_empty() {
        files.push("-".to_owned());
    }

    let settings = Settings::new(&matches);

    match wc(files, &settings) {
        Ok(()) => ( /* pass */ ),
        Err(e) => return e,
    }

    0
}

const CR: u8 = b'\r';
const LF: u8 = b'\n';
const SPACE: u8 = b' ';
const TAB: u8 = b'\t';
const SYN: u8 = 0x16 as u8;
const FF: u8 = 0x0C as u8;

#[inline(always)]
fn is_word_separator(byte: u8) -> bool {
    byte == SPACE || byte == TAB || byte == CR || byte == SYN || byte == FF
}

fn wc(files: Vec<String>, settings: &Settings) -> StdResult<(), i32> {
    let mut total_line_count: usize = 0;
    let mut total_word_count: usize = 0;
    let mut total_char_count: usize = 0;
    let mut total_byte_count: usize = 0;
    let mut total_longest_line_length: usize = 0;

    let mut results = vec![];
    let mut max_width: usize = 0;

    // we do not need to decode the byte stream if we're only counting bytes/newlines
    let decode_chars = settings.show_chars || settings.show_words || settings.show_max_line_length;

    for path in &files {
        let mut reader = open(&path[..])?;

        let mut line_count: usize = 0;
        let mut word_count: usize = 0;
        let mut byte_count: usize = 0;
        let mut char_count: usize = 0;
        let mut longest_line_length: usize = 0;
        let mut raw_line = Vec::new();

        // reading from a TTY seems to raise a condition on, rather than return Some(0) like a file.
        // hence the option wrapped in a result here
        while match reader.read_until(LF, &mut raw_line) {
            Ok(n) if n > 0 => true,
            Err(ref e) if !raw_line.is_empty() => {
                show_warning!("Error while reading {}: {}", path, e);
                !raw_line.is_empty()
            }
            _ => false,
        } {
            // GNU 'wc' only counts lines that end in LF as lines
            if *raw_line.last().unwrap() == LF {
                line_count += 1;
            }

            byte_count += raw_line.len();

            if decode_chars {
                // try and convert the bytes to UTF-8 first
                let current_char_count;
                match from_utf8(&raw_line[..]) {
                    Ok(line) => {
                        word_count += line.split_whitespace().count();
                        current_char_count = line.chars().count();
                    }
                    Err(..) => {
                        word_count += raw_line.split(|&x| is_word_separator(x)).count();
                        current_char_count = raw_line.iter().filter(|c| c.is_ascii()).count()
                    }
                }
                char_count += current_char_count;

                if current_char_count > longest_line_length {
                    // we subtract one here because `line.len()` includes the LF
                    // matches GNU 'wc' behavior
                    longest_line_length = current_char_count - 1;
                }
            }

            raw_line.truncate(0);
        }

        results.push(Result {
            title: path.clone(),
            bytes: byte_count,
            chars: char_count,
            lines: line_count,
            words: word_count,
            max_line_length: longest_line_length,
        });

        total_line_count += line_count;
        total_word_count += word_count;
        total_char_count += char_count;
        total_byte_count += byte_count;

        if longest_line_length > total_longest_line_length {
            total_longest_line_length = longest_line_length;
        }

        // used for formatting
        max_width = total_byte_count.to_string().len() + 1;
    }

    for result in &results {
        print_stats(settings, &result, max_width);
    }

    if files.len() > 1 {
        let result = Result {
            title: "total".to_owned(),
            bytes: total_byte_count,
            chars: total_char_count,
            lines: total_line_count,
            words: total_word_count,
            max_line_length: total_longest_line_length,
        };
        print_stats(settings, &result, max_width);
    }

    Ok(())
}

fn print_stats(settings: &Settings, result: &Result, max_width: usize) {
    if settings.show_lines {
        print!("{:1$}", result.lines, max_width);
    }
    if settings.show_words {
        print!("{:1$}", result.words, max_width);
    }
    if settings.show_bytes {
        print!("{:1$}", result.bytes, max_width);
    }
    if settings.show_chars {
        print!("{:1$}", result.chars, max_width);
    }
    if settings.show_max_line_length {
        print!("{:1$}", result.max_line_length, max_width);
    }

    if result.title != "-" {
        println!(" {}", result.title);
    } else {
        println!();
    }
}

fn open(path: &str) -> StdResult<BufReader<Box<dyn Read + 'static>>, i32> {
    if "-" == path {
        let reader = Box::new(stdin()) as Box<dyn Read>;
        return Ok(BufReader::new(reader));
    }

    let fpath = Path::new(path);
    if fpath.is_dir() {
        show_info!("{}: is a directory", path);
    }
    match File::open(&fpath) {
        Ok(fd) => {
            let reader = Box::new(fd) as Box<dyn Read>;
            Ok(BufReader::new(reader))
        }
        Err(e) => {
            show_error!("wc: {}: {}", path, e);
            Err(1)
        }
    }
}
<commit_msg>Update of the about description<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Boden Garman <bpgarman@gmail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

// spell-checker:ignore (ToDO) fpath

extern crate clap;
#[macro_use]
extern crate uucore;

use clap::{App, Arg, ArgMatches};

use std::fs::File;
use std::io::{stdin, BufRead, BufReader, Read};
use std::path::Path;
use std::result::Result as StdResult;
use std::str::from_utf8;

struct Settings {
    show_bytes: bool,
    show_chars: bool,
    show_lines: bool,
    show_words: bool,
    show_max_line_length: bool,
}

impl Settings {
    fn new(matches: &ArgMatches) -> Settings {
        let settings = Settings {
            show_bytes: matches.is_present(OPT_BYTES),
            show_chars: matches.is_present(OPT_CHAR),
            show_lines: matches.is_present(OPT_LINES),
            show_words: matches.is_present(OPT_WORDS),
            show_max_line_length: matches.is_present(OPT_MAX_LINE_LENGTH),
        };

        if settings.show_bytes
            || settings.show_chars
            || settings.show_lines
            || settings.show_words
            || settings.show_max_line_length
        {
            return settings;
        }

        Settings {
            show_bytes: true,
            show_chars: false,
            show_lines: true,
            show_words: true,
            show_max_line_length: false,
        }
    }
}

struct Result {
    title: String,
    bytes: usize,
    chars: usize,
    lines: usize,
    words: usize,
    max_line_length: usize,
}

static ABOUT: &str = "Display newline, word, and byte counts for each FILE, and a total line if
more than one FILE is specified.";
static VERSION: &str = env!("CARGO_PKG_VERSION");

static OPT_BYTES: &str = "bytes";
static OPT_CHAR: &str = "chars";
static OPT_LINES: &str = "lines";
static OPT_MAX_LINE_LENGTH: &str = "max-line-length";
static OPT_WORDS: &str = "words";

static ARG_FILES: &str = "files";

fn get_usage() -> String {
    format!(
        "{0} [OPTION]... [FILE]...
 With no FILE, or when FILE is -, read standard input.",
        executable!()
    )
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let usage = get_usage();

    let matches = App::new(executable!())
        .version(VERSION)
        .about(ABOUT)
        .usage(&usage[..])
        .arg(
            Arg::with_name(OPT_BYTES)
                .short("c")
                .long(OPT_BYTES)
                .help("print the byte counts"),
        )
        .arg(
            Arg::with_name(OPT_CHAR)
                .short("m")
                .long(OPT_CHAR)
                .help("print the character counts"),
        )
        .arg(
            Arg::with_name(OPT_LINES)
                .short("l")
                .long(OPT_LINES)
                .help("print the newline counts"),
        )
        .arg(
            Arg::with_name(OPT_MAX_LINE_LENGTH)
                .short("L")
                .long(OPT_MAX_LINE_LENGTH)
                .help("print the length of the longest line"),
        )
        .arg(
            Arg::with_name(OPT_WORDS)
                .short("w")
                .long(OPT_WORDS)
                .help("print the word counts"),
        )
        .arg(Arg::with_name(ARG_FILES).multiple(true).takes_value(true))
        .get_matches_from(args);

    let mut files: Vec<String> = matches
        .values_of(ARG_FILES)
        .map(|v| v.map(ToString::to_string).collect())
        .unwrap_or_default();

    if files.is_empty() {
        files.push("-".to_owned());
    }

    let settings = Settings::new(&matches);

    match wc(files, &settings) {
        Ok(()) => ( /* pass */ ),
        Err(e) => return e,
    }

    0
}

const CR: u8 = b'\r';
const LF: u8 = b'\n';
const SPACE: u8 = b' ';
const TAB: u8 = b'\t';
const SYN: u8 = 0x16 as u8;
const FF: u8 = 0x0C as u8;

#[inline(always)]
fn is_word_separator(byte: u8) -> bool {
    byte == SPACE || byte == TAB || byte == CR || byte == SYN || byte == FF
}

fn wc(files: Vec<String>, settings: &Settings) -> StdResult<(), i32> {
    let mut total_line_count: usize = 0;
    let mut total_word_count: usize = 0;
    let mut total_char_count: usize = 0;
    let mut total_byte_count: usize = 0;
    let mut total_longest_line_length: usize = 0;

    let mut results = vec![];
    let mut max_width: usize = 0;

    // we do not need to decode the byte stream if we're only counting bytes/newlines
    let decode_chars = settings.show_chars || settings.show_words || settings.show_max_line_length;

    for path in &files {
        let mut reader = open(&path[..])?;

        let mut line_count: usize = 0;
        let mut word_count: usize = 0;
        let mut byte_count: usize = 0;
        let mut char_count: usize = 0;
        let mut longest_line_length: usize = 0;
        let mut raw_line = Vec::new();

        // reading from a TTY seems to raise a condition on, rather than return Some(0) like a file.
        // hence the option wrapped in a result here
        while match reader.read_until(LF, &mut raw_line) {
            Ok(n) if n > 0 => true,
            Err(ref e) if !raw_line.is_empty() => {
                show_warning!("Error while reading {}: {}", path, e);
                !raw_line.is_empty()
            }
            _ => false,
        } {
            // GNU 'wc' only counts lines that end in LF as lines
            if *raw_line.last().unwrap() == LF {
                line_count += 1;
            }

            byte_count += raw_line.len();

            if decode_chars {
                // try and convert the bytes to UTF-8 first
                let current_char_count;
                match from_utf8(&raw_line[..]) {
                    Ok(line) => {
                        word_count += line.split_whitespace().count();
                        current_char_count = line.chars().count();
                    }
                    Err(..) => {
                        word_count += raw_line.split(|&x| is_word_separator(x)).count();
                        current_char_count = raw_line.iter().filter(|c| c.is_ascii()).count()
                    }
                }
                char_count += current_char_count;

                if current_char_count > longest_line_length {
                    // we subtract one here because `line.len()` includes the LF
                    // matches GNU 'wc' behavior
                    longest_line_length = current_char_count - 1;
                }
            }

            raw_line.truncate(0);
        }

        results.push(Result {
            title: path.clone(),
            bytes: byte_count,
            chars: char_count,
            lines: line_count,
            words: word_count,
            max_line_length: longest_line_length,
        });

        total_line_count += line_count;
        total_word_count += word_count;
        total_char_count += char_count;
        total_byte_count += byte_count;

        if longest_line_length > total_longest_line_length {
            total_longest_line_length = longest_line_length;
        }

        // used for formatting
        max_width = total_byte_count.to_string().len() + 1;
    }

    for result in &results {
        print_stats(settings, &result, max_width);
    }

    if files.len() > 1 {
        let result = Result {
            title: "total".to_owned(),
            bytes: total_byte_count,
            chars: total_char_count,
            lines: total_line_count,
            words: total_word_count,
            max_line_length: total_longest_line_length,
        };
        print_stats(settings, &result, max_width);
    }

    Ok(())
}

fn print_stats(settings: &Settings, result: &Result, max_width: usize) {
    if settings.show_lines {
        print!("{:1$}", result.lines, max_width);
    }
    if settings.show_words {
        print!("{:1$}", result.words, max_width);
    }
    if settings.show_bytes {
        print!("{:1$}", result.bytes, max_width);
    }
    if settings.show_chars {
        print!("{:1$}", result.chars, max_width);
    }
    if settings.show_max_line_length {
        print!("{:1$}", result.max_line_length, max_width);
    }

    if result.title != "-" {
        println!(" {}", result.title);
    } else {
        println!();
    }
}

fn open(path: &str) -> StdResult<BufReader<Box<dyn Read + 'static>>, i32> {
    if "-" == path {
        let reader = Box::new(stdin()) as Box<dyn Read>;
        return Ok(BufReader::new(reader));
    }

    let fpath = Path::new(path);
    if fpath.is_dir() {
        show_info!("{}: is a directory", path);
    }
    match File::open(&fpath) {
        Ok(fd) => {
            let reader = Box::new(fd) as Box<dyn Read>;
            Ok(BufReader::new(reader))
        }
        Err(e) => {
            show_error!("wc: {}: {}", path, e);
            Err(1)
        }
    }
}
<|endoftext|>


<commit_before>use futures_util::future::poll_fn;
use redis::{ConnectionAddr, ConnectionInfo, ErrorKind, RedisError, RedisResult, Value};
use std::collections::hash_map::{Entry, HashMap};
use std::future::Future;
use std::io::Result as IoResult;
use std::net::ToSocketAddrs;
use std::pin::Pin;
use std::task::Poll;
use tokio::io::{self, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::stream::Stream;
use tokio::sync::{broadcast, mpsc, oneshot};

#[cfg(unix)]
use tokio::net::UnixStream;

/// Read half of a Redis pubsub connection.
struct ReadHalf(io::BufReader<Box<dyn io::AsyncRead + Unpin + Send>>);
impl ReadHalf {
    /// Read a value from Redis.
    async fn read(&mut self) -> RedisResult<Value> {
        let mut decoder = combine::stream::Decoder::new();
        redis::parse_redis_value_async(&mut decoder, &mut self.0).await
    }
}

/// Write half of a Redis pubsub connection.
struct WriteHalf(Box<dyn io::AsyncWrite + Unpin + Send>);
impl WriteHalf {
    /// Write a command on the Redis connection.
    async fn write(&mut self, cmd: &[&[u8]]) -> IoResult<()> {
        let mut data = format!("*{}\r\n", cmd.len()).into_bytes();
        for part in cmd {
            data.append(&mut format!("${}\r\n", part.len()).into_bytes());
            data.extend_from_slice(part);
            data.extend_from_slice(b"\r\n");
        }
        self.0.write_all(&data).await
    }
}

/// Channel type used to receive pubsub messages.
pub type RecvChan = broadcast::Receiver<Vec<u8>>;

/// Channel type used to reply to `Cmd`.
type ReplyChan = oneshot::Sender<RecvChan>;

/// Command type sent to the connection loop.
struct Cmd {
    /// Channel to subscribe to.
    chan: Vec<u8>,
    /// Reply channel for the command.
    reply: ReplyChan,
}

/// Tracks an active subscription on the Redis server.
struct Sub {
    /// Channel sender used to notify all logical subscribers.
    tx: broadcast::Sender<Vec<u8>>,
    /// Reply channels that are awaiting confirmation.
    pending: Option<Vec<ReplyChan>>,
}

/// Polling events that can happen in the connection loop.
enum LoopEvent {
    Cmd(Cmd),
    CmdClosed,
    Interval,
    Read((RedisResult<Value>, ReadHalf)),
}

/// The Redis pubsub connection loop.
async fn conn_loop(mut rx: ReadHalf, mut tx: WriteHalf, mut cmd: mpsc::Receiver<Cmd>) {
    // Ping or cleanup subscriptions at an interval.
    let interval = tokio::time::interval(tokio::time::Duration::from_secs(20));
    tokio::pin!(interval);
    // Ignore the first (immediate) tick.
    interval.as_mut().tick().await;

    // Current read operation.
    // TODO: The redis crate has a ValueCodec, but doesn't expose it. This is a workaround.
    let mut read_fut: Pin<Box<dyn Future<Output = _> + Send>> = Box::pin(async move {
        let res = rx.read().await;
        (res, rx)
    });

    // Map of active subscriptions by channel name.
    let mut subs: HashMap<Vec<u8>, Sub> = HashMap::new();

    loop {
        // This specifically prioritizes processing commands over receiving.
        match poll_fn(|cx| {
            if let Poll::Ready(res) = cmd.poll_recv(cx) {
                match res {
                    Some(cmd) => Poll::Ready(LoopEvent::Cmd(cmd)),
                    None => Poll::Ready(LoopEvent::CmdClosed),
                }
            } else if let Poll::Ready(_) = interval.as_mut().poll_next(cx) {
                Poll::Ready(LoopEvent::Interval)
            } else if let Poll::Ready(res) = read_fut.as_mut().poll(cx) {
                Poll::Ready(LoopEvent::Read(res))
            } else {
                Poll::Pending
            }
        })
        .await
        {
            LoopEvent::Cmd(Cmd { chan, reply }) => match subs.entry(chan.clone()) {
                // If already subscribed, reply with a broadcast channel immediately. Otherwise,
                // add the reply channel to `pending`, and send the Redis subscribe command if
                // necessary.
                Entry::Occupied(mut entry) => {
                    let sub = entry.get_mut();
                    if let Some(ref mut pending) = sub.pending {
                        pending.push(reply);
                    } else {
                        let _ = reply.send(sub.tx.subscribe());
                    }
                }
                Entry::Vacant(entry) => {
                    entry.insert(Sub {
                        tx: broadcast::channel(8).0,
                        pending: Some(vec![reply]),
                    });
                    tx.write(&[b"SUBSCRIBE", &chan])
                        .await
                        .expect("Failed to send subscribe command to Redis");
                }
            },
            LoopEvent::CmdClosed => {
                // TODO: Stop reading from the command channel.
                // This, plus an empty `subs`, means we can exit.
                unimplemented!();
            }
            LoopEvent::Interval => {
                // Unsubscribe from channels that no longer have subscribers, or send a ping.
                let to_unsub: Vec<Vec<u8>> = subs
                    .iter()
                    .filter_map(|(chan, sub)| {
                        if sub.pending.is_none() && sub.tx.receiver_count() == 0 {
                            Some(chan.clone())
                        } else {
                            None
                        }
                    })
                    .collect();
                if to_unsub.is_empty() {
                    tx.write(&[b"PING"])
                        .await
                        .expect("Failed to send ping command to Redis");
                } else {
                    for chan in &to_unsub {
                        subs.remove(chan);
                    }
                    let mut unsub_cmd: Vec<&[u8]> = vec![b"UNSUBSCRIBE"];
                    unsub_cmd.extend(to_unsub.iter().map(|chan| &chan[..]));
                    tx.write(&unsub_cmd)
                        .await
                        .expect("Failed to send unsubscribe command to Redis");
                }
            }
            LoopEvent::Read((res, mut rx)) => {
                read_fut = Box::pin(async move {
                    let res = rx.read().await;
                    (res, rx)
                });
                let value = res.expect("Failed to read from Redis");
                let vec = match value {
                    // Note: If we have no subscriptions at all, we receive pongs as regular
                    // replies instead of events.
                    Value::Status(status) if status == "PONG" => continue,
                    Value::Bulk(ref vec) if vec.len() >= 2 => vec,
                    _ => panic!("Unexpected value from Redis: {:?}", value),
                };
                match (&vec[0], &vec[1], vec.get(2)) {
                    // Handle a message event by sending on the broadcast channel.
                    (
                        &Value::Data(ref ev),
                        &Value::Data(ref chan),
                        Some(&Value::Data(ref data)),
                    ) if ev == b"message" => {
                        if let Some(ref sub) = subs.get(&chan[..]) {
                            let _ = sub.tx.send(data.to_vec());
                        }
                    }
                    // Handle subscription confirmation by sending out pending replies.
                    (&Value::Data(ref ev), &Value::Data(ref chan), _) if ev == b"subscribe" => {
                        if let Some(ref mut sub) = subs.get_mut(&chan[..]) {
                            if let Some(pending) = sub.pending.take() {
                                for reply in pending {
                                    let _ = reply.send(sub.tx.subscribe());
                                }
                            }
                        }
                    }
                    // Some other events are ok, but we do nothing with them.
                    (&Value::Data(ref ev), _, _) if ev == b"unsubscribe" || ev == b"pong" => {}
                    _ => panic!("Unexpected value from Redis: {:?}", value),
                }
            }
        }
    }
}

/// A Subscriber can be used to subscribe to Redis pubsub channels.
///
/// This struct can be cheaply cloned. It is simply a client of the connection loop which is
/// running in another task.
#[derive(Clone)]
pub struct Subscriber {
    cmd: mpsc::Sender<Cmd>,
}

impl Subscriber {
    /// Subscribe to a channel.
    ///
    /// This function does not complete until the server has confirmed the subscription.
    pub async fn subscribe(&mut self, chan: Vec<u8>) -> broadcast::Receiver<Vec<u8>> {
        let (reply_tx, reply_rx) = oneshot::channel();
        let cmd = Cmd {
            chan,
            reply: reply_tx,
        };
        if self.cmd.send(cmd).await.is_ok() {
            if let Ok(rx) = reply_rx.await {
                return rx;
            }
        }
        panic!("Tried to subscribe on closed pubsub connection");
    }
}

/// Make a pubsub connection to Redis.
pub async fn connect(info: &ConnectionInfo) -> RedisResult<Subscriber> {
    // Note: This code is borrowed from the redis crate.
    let (rx, tx): (
        Box<dyn io::AsyncRead + Unpin + Send>,
        Box<dyn io::AsyncWrite + Unpin + Send>,
    ) = match *info.addr {
        ConnectionAddr::Tcp(ref host, port) => {
            let socket_addr = {
                let mut socket_addrs = (&host[..], port).to_socket_addrs()?;
                match socket_addrs.next() {
                    Some(socket_addr) => socket_addr,
                    None => {
                        return Err(RedisError::from((
                            ErrorKind::InvalidClientConfig,
                            "No address found for host",
                        )));
                    }
                }
            };

            let (rx, tx) = io::split(TcpStream::connect(&socket_addr).await?);
            (Box::new(rx), Box::new(tx))
        }

        ConnectionAddr::TcpTls { .. } => {
            return Err(RedisError::from((ErrorKind::InvalidClientConfig, "TLS connections not yet supported")))
        }

        #[cfg(unix)]
        ConnectionAddr::Unix(ref path) => {
            let (rx, tx) = io::split(UnixStream::connect(path).await?);
            (Box::new(rx), Box::new(tx))
        }

        #[cfg(not(unix))]
        ConnectionAddr::Unix(_) => {
            return Err(RedisError::from((
                ErrorKind::InvalidClientConfig,
                "Cannot connect to unix sockets \
                 on this platform",
            )))
        }
    };

    let mut rx = ReadHalf(io::BufReader::new(rx));
    let mut tx = WriteHalf(tx);

    if let Some(ref passwd) = info.passwd {
        tx.write(&[b"AUTH", passwd.as_bytes()]).await?;
        match rx.read().await {
            Ok(Value::Okay) => (),
            _ => {
                return Err((
                    ErrorKind::AuthenticationFailed,
                    "Password authentication failed",
                )
                    .into());
            }
        }
    }

    // Note: Pubsub ignores database ID, so we don't need to send `SELECT`.

    let (cmd_tx, cmd_rx) = mpsc::channel(8);
    tokio::spawn(conn_loop(rx, tx, cmd_rx));
    Ok(Subscriber { cmd: cmd_tx })
}
<commit_msg>cargo fmt<commit_after>use futures_util::future::poll_fn;
use redis::{ConnectionAddr, ConnectionInfo, ErrorKind, RedisError, RedisResult, Value};
use std::collections::hash_map::{Entry, HashMap};
use std::future::Future;
use std::io::Result as IoResult;
use std::net::ToSocketAddrs;
use std::pin::Pin;
use std::task::Poll;
use tokio::io::{self, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::stream::Stream;
use tokio::sync::{broadcast, mpsc, oneshot};

#[cfg(unix)]
use tokio::net::UnixStream;

/// Read half of a Redis pubsub connection.
struct ReadHalf(io::BufReader<Box<dyn io::AsyncRead + Unpin + Send>>);
impl ReadHalf {
    /// Read a value from Redis.
    async fn read(&mut self) -> RedisResult<Value> {
        let mut decoder = combine::stream::Decoder::new();
        redis::parse_redis_value_async(&mut decoder, &mut self.0).await
    }
}

/// Write half of a Redis pubsub connection.
struct WriteHalf(Box<dyn io::AsyncWrite + Unpin + Send>);
impl WriteHalf {
    /// Write a command on the Redis connection.
    async fn write(&mut self, cmd: &[&[u8]]) -> IoResult<()> {
        let mut data = format!("*{}\r\n", cmd.len()).into_bytes();
        for part in cmd {
            data.append(&mut format!("${}\r\n", part.len()).into_bytes());
            data.extend_from_slice(part);
            data.extend_from_slice(b"\r\n");
        }
        self.0.write_all(&data).await
    }
}

/// Channel type used to receive pubsub messages.
pub type RecvChan = broadcast::Receiver<Vec<u8>>;

/// Channel type used to reply to `Cmd`.
type ReplyChan = oneshot::Sender<RecvChan>;

/// Command type sent to the connection loop.
struct Cmd {
    /// Channel to subscribe to.
    chan: Vec<u8>,
    /// Reply channel for the command.
    reply: ReplyChan,
}

/// Tracks an active subscription on the Redis server.
struct Sub {
    /// Channel sender used to notify all logical subscribers.
    tx: broadcast::Sender<Vec<u8>>,
    /// Reply channels that are awaiting confirmation.
    pending: Option<Vec<ReplyChan>>,
}

/// Polling events that can happen in the connection loop.
enum LoopEvent {
    Cmd(Cmd),
    CmdClosed,
    Interval,
    Read((RedisResult<Value>, ReadHalf)),
}

/// The Redis pubsub connection loop.
async fn conn_loop(mut rx: ReadHalf, mut tx: WriteHalf, mut cmd: mpsc::Receiver<Cmd>) {
    // Ping or cleanup subscriptions at an interval.
    let interval = tokio::time::interval(tokio::time::Duration::from_secs(20));
    tokio::pin!(interval);
    // Ignore the first (immediate) tick.
    interval.as_mut().tick().await;

    // Current read operation.
    // TODO: The redis crate has a ValueCodec, but doesn't expose it. This is a workaround.
    let mut read_fut: Pin<Box<dyn Future<Output = _> + Send>> = Box::pin(async move {
        let res = rx.read().await;
        (res, rx)
    });

    // Map of active subscriptions by channel name.
    let mut subs: HashMap<Vec<u8>, Sub> = HashMap::new();

    loop {
        // This specifically prioritizes processing commands over receiving.
        match poll_fn(|cx| {
            if let Poll::Ready(res) = cmd.poll_recv(cx) {
                match res {
                    Some(cmd) => Poll::Ready(LoopEvent::Cmd(cmd)),
                    None => Poll::Ready(LoopEvent::CmdClosed),
                }
            } else if let Poll::Ready(_) = interval.as_mut().poll_next(cx) {
                Poll::Ready(LoopEvent::Interval)
            } else if let Poll::Ready(res) = read_fut.as_mut().poll(cx) {
                Poll::Ready(LoopEvent::Read(res))
            } else {
                Poll::Pending
            }
        })
        .await
        {
            LoopEvent::Cmd(Cmd { chan, reply }) => match subs.entry(chan.clone()) {
                // If already subscribed, reply with a broadcast channel immediately. Otherwise,
                // add the reply channel to `pending`, and send the Redis subscribe command if
                // necessary.
                Entry::Occupied(mut entry) => {
                    let sub = entry.get_mut();
                    if let Some(ref mut pending) = sub.pending {
                        pending.push(reply);
                    } else {
                        let _ = reply.send(sub.tx.subscribe());
                    }
                }
                Entry::Vacant(entry) => {
                    entry.insert(Sub {
                        tx: broadcast::channel(8).0,
                        pending: Some(vec![reply]),
                    });
                    tx.write(&[b"SUBSCRIBE", &chan])
                        .await
                        .expect("Failed to send subscribe command to Redis");
                }
            },
            LoopEvent::CmdClosed => {
                // TODO: Stop reading from the command channel.
                // This, plus an empty `subs`, means we can exit.
                unimplemented!();
            }
            LoopEvent::Interval => {
                // Unsubscribe from channels that no longer have subscribers, or send a ping.
                let to_unsub: Vec<Vec<u8>> = subs
                    .iter()
                    .filter_map(|(chan, sub)| {
                        if sub.pending.is_none() && sub.tx.receiver_count() == 0 {
                            Some(chan.clone())
                        } else {
                            None
                        }
                    })
                    .collect();
                if to_unsub.is_empty() {
                    tx.write(&[b"PING"])
                        .await
                        .expect("Failed to send ping command to Redis");
                } else {
                    for chan in &to_unsub {
                        subs.remove(chan);
                    }
                    let mut unsub_cmd: Vec<&[u8]> = vec![b"UNSUBSCRIBE"];
                    unsub_cmd.extend(to_unsub.iter().map(|chan| &chan[..]));
                    tx.write(&unsub_cmd)
                        .await
                        .expect("Failed to send unsubscribe command to Redis");
                }
            }
            LoopEvent::Read((res, mut rx)) => {
                read_fut = Box::pin(async move {
                    let res = rx.read().await;
                    (res, rx)
                });
                let value = res.expect("Failed to read from Redis");
                let vec = match value {
                    // Note: If we have no subscriptions at all, we receive pongs as regular
                    // replies instead of events.
                    Value::Status(status) if status == "PONG" => continue,
                    Value::Bulk(ref vec) if vec.len() >= 2 => vec,
                    _ => panic!("Unexpected value from Redis: {:?}", value),
                };
                match (&vec[0], &vec[1], vec.get(2)) {
                    // Handle a message event by sending on the broadcast channel.
                    (
                        &Value::Data(ref ev),
                        &Value::Data(ref chan),
                        Some(&Value::Data(ref data)),
                    ) if ev == b"message" => {
                        if let Some(ref sub) = subs.get(&chan[..]) {
                            let _ = sub.tx.send(data.to_vec());
                        }
                    }
                    // Handle subscription confirmation by sending out pending replies.
                    (&Value::Data(ref ev), &Value::Data(ref chan), _) if ev == b"subscribe" => {
                        if let Some(ref mut sub) = subs.get_mut(&chan[..]) {
                            if let Some(pending) = sub.pending.take() {
                                for reply in pending {
                                    let _ = reply.send(sub.tx.subscribe());
                                }
                            }
                        }
                    }
                    // Some other events are ok, but we do nothing with them.
                    (&Value::Data(ref ev), _, _) if ev == b"unsubscribe" || ev == b"pong" => {}
                    _ => panic!("Unexpected value from Redis: {:?}", value),
                }
            }
        }
    }
}

/// A Subscriber can be used to subscribe to Redis pubsub channels.
///
/// This struct can be cheaply cloned. It is simply a client of the connection loop which is
/// running in another task.
#[derive(Clone)]
pub struct Subscriber {
    cmd: mpsc::Sender<Cmd>,
}

impl Subscriber {
    /// Subscribe to a channel.
    ///
    /// This function does not complete until the server has confirmed the subscription.
    pub async fn subscribe(&mut self, chan: Vec<u8>) -> broadcast::Receiver<Vec<u8>> {
        let (reply_tx, reply_rx) = oneshot::channel();
        let cmd = Cmd {
            chan,
            reply: reply_tx,
        };
        if self.cmd.send(cmd).await.is_ok() {
            if let Ok(rx) = reply_rx.await {
                return rx;
            }
        }
        panic!("Tried to subscribe on closed pubsub connection");
    }
}

/// Make a pubsub connection to Redis.
pub async fn connect(info: &ConnectionInfo) -> RedisResult<Subscriber> {
    // Note: This code is borrowed from the redis crate.
    let (rx, tx): (
        Box<dyn io::AsyncRead + Unpin + Send>,
        Box<dyn io::AsyncWrite + Unpin + Send>,
    ) = match *info.addr {
        ConnectionAddr::Tcp(ref host, port) => {
            let socket_addr = {
                let mut socket_addrs = (&host[..], port).to_socket_addrs()?;
                match socket_addrs.next() {
                    Some(socket_addr) => socket_addr,
                    None => {
                        return Err(RedisError::from((
                            ErrorKind::InvalidClientConfig,
                            "No address found for host",
                        )));
                    }
                }
            };

            let (rx, tx) = io::split(TcpStream::connect(&socket_addr).await?);
            (Box::new(rx), Box::new(tx))
        }

        ConnectionAddr::TcpTls { .. } => {
            return Err(RedisError::from((
                ErrorKind::InvalidClientConfig,
                "TLS connections not yet supported",
            )))
        }

        #[cfg(unix)]
        ConnectionAddr::Unix(ref path) => {
            let (rx, tx) = io::split(UnixStream::connect(path).await?);
            (Box::new(rx), Box::new(tx))
        }

        #[cfg(not(unix))]
        ConnectionAddr::Unix(_) => {
            return Err(RedisError::from((
                ErrorKind::InvalidClientConfig,
                "Cannot connect to unix sockets \
                 on this platform",
            )))
        }
    };

    let mut rx = ReadHalf(io::BufReader::new(rx));
    let mut tx = WriteHalf(tx);

    if let Some(ref passwd) = info.passwd {
        tx.write(&[b"AUTH", passwd.as_bytes()]).await?;
        match rx.read().await {
            Ok(Value::Okay) => (),
            _ => {
                return Err((
                    ErrorKind::AuthenticationFailed,
                    "Password authentication failed",
                )
                    .into());
            }
        }
    }

    // Note: Pubsub ignores database ID, so we don't need to send `SELECT`.

    let (cmd_tx, cmd_rx) = mpsc::channel(8);
    tokio::spawn(conn_loop(rx, tx, cmd_rx));
    Ok(Subscriber { cmd: cmd_tx })
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

<%namespace name="helpers" file="/helpers.mako.rs" />
<% from data import ALL_SIDES, maybe_moz_logical_alias %>
<% data.new_style_struct("Padding", inherited=False) %>

% for side in ALL_SIDES:
    <%
        spec = "https://drafts.csswg.org/css-box/#propdef-padding-%s" % side[0]
        if side[1]:
            spec = "https://drafts.csswg.org/css-logical-props/#propdef-padding-%s" % side[1]
    %>
    ${helpers.predefined_type("padding-%s" % side[0], "LengthOrPercentage",
                               "computed::LengthOrPercentage::Length(Au(0))",
                               "parse_non_negative",
                               alias=maybe_moz_logical_alias(product, side, "-moz-padding-%s"),
                               needs_context=False,
                               animatable=True,
                               logical = side[1],
                               spec = spec)}
% endfor
<commit_msg>style: Fix alignment of the padding properties.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

<%namespace name="helpers" file="/helpers.mako.rs" />
<% from data import ALL_SIDES, maybe_moz_logical_alias %>
<% data.new_style_struct("Padding", inherited=False) %>

% for side in ALL_SIDES:
    <%
        spec = "https://drafts.csswg.org/css-box/#propdef-padding-%s" % side[0]
        if side[1]:
            spec = "https://drafts.csswg.org/css-logical-props/#propdef-padding-%s" % side[1]
    %>
    ${helpers.predefined_type("padding-%s" % side[0], "LengthOrPercentage",
                              "computed::LengthOrPercentage::Length(Au(0))",
                              "parse_non_negative",
                              alias=maybe_moz_logical_alias(product, side, "-moz-padding-%s"),
                              needs_context=False,
                              animatable=True,
                              logical = side[1],
                              spec = spec)}
% endfor
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

<%namespace name="helpers" file="/helpers.mako.rs" />
<% from data import ALL_SIDES, maybe_moz_logical_alias %>
<% data.new_style_struct("Margin", inherited=False) %>

% for side in ALL_SIDES:
    <%
        spec = "https://drafts.csswg.org/css-box/#propdef-margin-%s" % side[0]
        if side[1]:
            spec = "https://drafts.csswg.org/css-logical-props/#propdef-margin-%s" % side[1]
    %>
    ${helpers.predefined_type(
        "margin-%s" % side[0],
        "LengthPercentageOrAuto",
        "computed::LengthPercentageOrAuto::zero()",
        alias=maybe_moz_logical_alias(product, side, "-moz-margin-%s"),
        allow_quirks="No" if side[1] else "Yes",
        animation_value_type="ComputedValue",
        logical=side[1],
        logical_group="margin",
        spec=spec,
        flags="APPLIES_TO_FIRST_LETTER GETCS_NEEDS_LAYOUT_FLUSH",
        allowed_in_page_rule=True,
        servo_restyle_damage="reflow"
    )}
% endfor

% for side in ALL_SIDES:
    ${helpers.predefined_type(
        "scroll-margin-%s" % side[0],
        "Length",
        "computed::Length::zero()",
        products="gecko",
        gecko_pref="layout.css.scroll-snap-v1.enabled",
        logical=side[1],
        logical_group="scroll-margin",
        spec="https://drafts.csswg.org/css-scroll-snap-1/#propdef-scroll-margin-%s" % side[0],
        flags="APPLIES_TO_FIRST_LETTER APPLIES_TO_FIRST_LINE",
        animation_value_type="ComputedValue",
    )}
% endfor
<commit_msg>style: scroll-margin properties shouldn't apply to first-line / first-letter.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

<%namespace name="helpers" file="/helpers.mako.rs" />
<% from data import ALL_SIDES, maybe_moz_logical_alias %>
<% data.new_style_struct("Margin", inherited=False) %>

% for side in ALL_SIDES:
    <%
        spec = "https://drafts.csswg.org/css-box/#propdef-margin-%s" % side[0]
        if side[1]:
            spec = "https://drafts.csswg.org/css-logical-props/#propdef-margin-%s" % side[1]
    %>
    ${helpers.predefined_type(
        "margin-%s" % side[0],
        "LengthPercentageOrAuto",
        "computed::LengthPercentageOrAuto::zero()",
        alias=maybe_moz_logical_alias(product, side, "-moz-margin-%s"),
        allow_quirks="No" if side[1] else "Yes",
        animation_value_type="ComputedValue",
        logical=side[1],
        logical_group="margin",
        spec=spec,
        flags="APPLIES_TO_FIRST_LETTER GETCS_NEEDS_LAYOUT_FLUSH",
        allowed_in_page_rule=True,
        servo_restyle_damage="reflow"
    )}
% endfor

% for side in ALL_SIDES:
    ${helpers.predefined_type(
        "scroll-margin-%s" % side[0],
        "Length",
        "computed::Length::zero()",
        products="gecko",
        gecko_pref="layout.css.scroll-snap-v1.enabled",
        logical=side[1],
        logical_group="scroll-margin",
        spec="https://drafts.csswg.org/css-scroll-snap-1/#propdef-scroll-margin-%s" % side[0],
        animation_value_type="ComputedValue",
    )}
% endfor
<|endoftext|>


<commit_before>use regex::Regex;
use chrono::NaiveDate;
use reqwest;
use std::io::Read;
use select::document::Document;
use select::predicate::{Name, Attr};
use rustc_serialize::json::Json;
use errors::*;

// Some formats borrowed from https://github.com/amir/article-date-extractor
static FMTS: &'static [&str] = &["%A, %B %e, %Y", "%Y-%m-%dT%H:%M:%S%:z", "/%Y/%m/%d/", "/%Y/%d/%m/", "%Y-%m-%d", "%B %e, %Y", "%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%SZ", "%B %k, %Y, %H:%M %p", "%Y-%m-%d %H:%M:%S.000000"];

// Use lazy_static to ensure we only compile the regex once
lazy_static! {
    // Regex by Newspaper3k  - https://github.com/codelucas/newspaper/blob/master/newspaper/urls.py
    static ref RE: Regex = Regex::new(r"([\./\-_]{0,1}(19|20)\d{2})[\./\-_]{0,1}(([0-3]{0,1}[0-9][\./\-_])|(\w{3,5}[\./\-_]))([0-3]{0,1}[0-9][\./\-]{0,1})").unwrap();
}

// Parse the date, trying out each format
fn parse_date(input: &str) -> Result<NaiveDate> {
    let mut result: Result<NaiveDate> = Err("None of the formats matched the date".into());

    'outer: for fmt in FMTS {
        if let Ok(v) = NaiveDate::parse_from_str(input, fmt) {
            { result = Ok(v);
              break 'outer; }
        }
    }

    result
}

// Extract date from a URL
fn extract_from_url(url: &str) -> Option<String> {
    if let Some(val) = RE.find(url) {
        return Some(val.as_str().to_string())
    } else {
        return None
    }
}

// Extract date from JSON-LD
fn extract_from_ldjson<'a>(html: &'a Document) -> Option<String> {
    let mut json_date: Option<String> = None;
    let mut ldjson = String::new();
    if let Some(ldj) = html.find(Attr("type", "application/ld+json")).next() {
        ldjson = ldj.text();
    }

    let mut decoded_ldjson = Json::from_str("{}").unwrap();

    match Json::from_str(ldjson.as_str()) {
        Ok(v) => decoded_ldjson = v,
        _ => return None,
    }

    if let Some(date_published) = decoded_ldjson.search("datePublished") {
        if let Some(date) = date_published.as_string() {
            json_date = Some(date.to_string())
        }
    } else if let Some(date_created) = decoded_ldjson.search("dateCreated") {
        if let Some(date) = date_created.as_string() {
            json_date = Some(date.to_string())
        }
    }

    json_date
}

// Extract date from meta tags
fn extract_from_meta<'a>(html: &'a Document) -> Option<String> {
    let mut meta_date: Option<String> = None;

    'outer: for meta in html.find(Name("meta")) {
        let meta_name     = meta.attr("name").unwrap_or("").to_lowercase();
        let item_prop     = meta.attr("itemprop").unwrap_or("").to_lowercase();
        let http_equiv    = meta.attr("http-equiv").unwrap_or("").to_lowercase();
        let meta_property = meta.attr("property").unwrap_or("").to_lowercase();

        match meta_name.as_ref() {
            "pubdate"               | "publishdate"                  | "timestamp"       |
            "dc.date.issued"        | "date"                         | "sailthru.date"   |
            "article.published"     | "published-date"               | "article.created" |
            "article_date_original" | "cxenseparse:recs:publishtime" | "date_published"  => { if let Some(ct) = meta.attr("content") {
                                                                                meta_date = Some(ct.trim().to_string())
                                                                                }
                                                                                break 'outer; },
            _ => {},
        }

        match item_prop.as_ref() {
            "datepublished" | "datecreated" => { if let Some(ct) = meta.attr("content") {
                                                   meta_date = Some(ct.trim().to_string())
                                                 }
                                                 break 'outer; },
            _ => {},
        }

        match http_equiv.as_ref() {
            "date" =>  { if let Some(ct) = meta.attr("content") {
                           meta_date = Some(ct.trim().to_string())
                         }
                         break 'outer; },
            _ => {},
        }

        match meta_property.as_ref() {
            "article:published_time" | "bt:pubdate" => { if let Some(ct) = meta.attr("content") {
                                                           meta_date = Some(ct.trim().to_string())
                                                         }
                                                         break 'outer; },
            "og:image"                              => { if let Some(url) = meta.attr("content") {
                                                           meta_date = extract_from_url(url.trim())
                                                         }
                                                         break 'outer; },

            _ => {},
        }

    }

    meta_date
}

// Extract from html tags
fn extract_from_html_tag<'a>(html: &'a Document) -> Option<String> {
    let mut date: Option<String> = None;

    for time in html.find(Name("time")) {
        if let Some(dt) = time.attr("datetime") {
            date = Some(dt.to_string())
        } else if let Some("timestamp") = time.attr("class") {
            date = Some(time.text())
        }
    }

    if date.is_none() {
        for tag in html.find(Name("meta")) {
            if let Some("datePublished") = tag.attr("itemprop") {
                if let Some(v) = tag.attr("content") {
                    date = Some(v.to_string())
                } else {
                    date = Some(tag.text())
                }
            }
        }
    }

    date
}

// Try to extract the date by using each function one by one
pub fn extract_article_published_date(link: &str, html: Option<String>) -> Result<NaiveDate> {
    let mut body: String = String::new();
    let mut parsed_body: Option<Document> = None;

    if let Some(v) = extract_from_url(link) {
        return parse_date(v.as_str())
    }

    if html.is_none() {
        if let Ok(mut response) = reqwest::get(link) {
            response.read_to_string(&mut body).unwrap();
            let doc = Document::from(body.as_str());
            parsed_body = Some(doc);
        } else {
            return Err("Couldn't open the link".into())
        }
    } else {
        parsed_body = Some(Document::from(html.unwrap().as_str()))
    }

    if let Some(v) = extract_from_url(link) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_ldjson(parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_meta(parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_html_tag(parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else {
        return Err("Couldn't find the date to parse".into())
    }
}

// Unit tests
#[cfg(test)]
mod test {
    use super::extract_from_url;
    use super::parse_date;
    use super::extract_from_meta;
    use super::extract_from_ldjson;
    use super::extract_from_html_tag;
    use super::extract_article_published_date;
    use chrono::NaiveDate;
    use reqwest;
    use std::string::String;
    use std::io::Read;
    use select::document::Document;

    #[test]
    fn parsing_date() {
        assert_eq!(NaiveDate::from_ymd(2015,11,30), parse_date("/2015/11/30/").unwrap());
        assert_eq!(NaiveDate::from_ymd(2015,11,30), parse_date("/2015/30/11/").unwrap());

        assert!(parse_date("bad_format").is_err());
    }

    #[test]
    fn extracting_from_url() {
        let link = "http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html";
        assert_eq!(Some("/2015/11/28/".to_string()), extract_from_url(link));

        let link = "";
        assert_eq!(None, extract_from_url(link));
    }

    #[test]
    fn extracting_from_meta() {
        let mut response = reqwest::get("https://techcrunch.com/2015/11/30/atlassian-share-price/").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();
        let document = Document::from(body.as_str());

        assert_eq!(Some(("2015-11-30 23:50:48".to_string())), extract_from_meta(&document));
    }

    #[test]
    fn extracting_from_ldjson() {
        let mut response = reqwest::get("https://techcrunch.com/2015/11/30/atlassian-share-price/").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();
        let document = Document::from(body.as_str());

        assert_eq!(Some("2015-12-01T07:50:48Z".to_string()), extract_from_ldjson(&document));
    }

    #[test]
    fn extracting_from_html_tag() {
        let mut response = reqwest::get("http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();
        let document = Document::from(body.as_str());

        assert_eq!(Some("2015-11-29T00:44:59Z".to_string()), extract_from_html_tag(&document));
    }

    #[test]
    fn extracting_article_published_date() {
        let link = "http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html";
        let mut response = reqwest::get("http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();

        assert_eq!(NaiveDate::from_ymd(2015,11,28), extract_article_published_date(&link, None).unwrap());
        assert_eq!(NaiveDate::from_ymd(2015,11,28), extract_article_published_date(&link, Some(body)).unwrap());

        assert!((extract_article_published_date("", None)).is_err());
    }
}
<commit_msg>Improve extract_from_html_tag<commit_after>use regex::Regex;
use chrono::NaiveDate;
use reqwest;
use std::io::Read;
use select::document::Document;
use select::predicate::{Name, Attr};
use rustc_serialize::json::Json;
use errors::*;

// Some formats borrowed from https://github.com/amir/article-date-extractor
static FMTS: &'static [&str] = &["%A, %B %e, %Y", "%Y-%m-%dT%H:%M:%S%:z", "/%Y/%m/%d/", "/%Y/%d/%m/", "%Y-%m-%d", "%B %e, %Y", "%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%SZ", "%B %k, %Y, %H:%M %p", "%Y-%m-%d %H:%M:%S.000000"];

// Use lazy_static to ensure we only compile the regex once
lazy_static! {
    // Regex by Newspaper3k  - https://github.com/codelucas/newspaper/blob/master/newspaper/urls.py
    static ref RE: Regex = Regex::new(r"([\./\-_]{0,1}(19|20)\d{2})[\./\-_]{0,1}(([0-3]{0,1}[0-9][\./\-_])|(\w{3,5}[\./\-_]))([0-3]{0,1}[0-9][\./\-]{0,1})").unwrap();
}

// Parse the date, trying out each format
fn parse_date(input: &str) -> Result<NaiveDate> {
    let mut result: Result<NaiveDate> = Err("None of the formats matched the date".into());

    'outer: for fmt in FMTS {
        if let Ok(v) = NaiveDate::parse_from_str(input, fmt) {
            { result = Ok(v);
              break 'outer; }
        }
    }

    result
}

// Extract date from a URL
fn extract_from_url(url: &str) -> Option<String> {
    if let Some(val) = RE.find(url) {
        return Some(val.as_str().to_string())
    } else {
        return None
    }
}

// Extract date from JSON-LD
fn extract_from_ldjson<'a>(html: &'a Document) -> Option<String> {
    let mut json_date: Option<String> = None;
    let mut ldjson = String::new();
    if let Some(ldj) = html.find(Attr("type", "application/ld+json")).next() {
        ldjson = ldj.text();
    }

    let mut decoded_ldjson = Json::from_str("{}").unwrap();

    match Json::from_str(ldjson.as_str()) {
        Ok(v) => decoded_ldjson = v,
        _ => return None,
    }

    if let Some(date_published) = decoded_ldjson.search("datePublished") {
        if let Some(date) = date_published.as_string() {
            json_date = Some(date.to_string())
        }
    } else if let Some(date_created) = decoded_ldjson.search("dateCreated") {
        if let Some(date) = date_created.as_string() {
            json_date = Some(date.to_string())
        }
    }

    json_date
}

// Extract date from meta tags
fn extract_from_meta<'a>(html: &'a Document) -> Option<String> {
    let mut meta_date: Option<String> = None;

    'outer: for meta in html.find(Name("meta")) {
        let meta_name     = meta.attr("name").unwrap_or("").to_lowercase();
        let item_prop     = meta.attr("itemprop").unwrap_or("").to_lowercase();
        let http_equiv    = meta.attr("http-equiv").unwrap_or("").to_lowercase();
        let meta_property = meta.attr("property").unwrap_or("").to_lowercase();

        match meta_name.as_ref() {
            "pubdate"               | "publishdate"                  | "timestamp"       |
            "dc.date.issued"        | "date"                         | "sailthru.date"   |
            "article.published"     | "published-date"               | "article.created" |
            "article_date_original" | "cxenseparse:recs:publishtime" | "date_published"  => { if let Some(ct) = meta.attr("content") {
                                                                                meta_date = Some(ct.trim().to_string())
                                                                                }
                                                                                break 'outer; },
            _ => {},
        }

        match item_prop.as_ref() {
            "datepublished" | "datecreated" => { if let Some(ct) = meta.attr("content") {
                                                   meta_date = Some(ct.trim().to_string())
                                                 }
                                                 break 'outer; },
            _ => {},
        }

        match http_equiv.as_ref() {
            "date" =>  { if let Some(ct) = meta.attr("content") {
                           meta_date = Some(ct.trim().to_string())
                         }
                         break 'outer; },
            _ => {},
        }

        match meta_property.as_ref() {
            "article:published_time" | "bt:pubdate" => { if let Some(ct) = meta.attr("content") {
                                                           meta_date = Some(ct.trim().to_string())
                                                         }
                                                         break 'outer; },
            "og:image"                              => { if let Some(url) = meta.attr("content") {
                                                           meta_date = extract_from_url(url.trim())
                                                         }
                                                         break 'outer; },

            _ => {},
        }

    }

    meta_date
}

// Extract from html tags
fn extract_from_html_tag<'a>(html: &'a Document) -> Option<String> {
    lazy_static! {
        static ref TAG_RE: Regex = Regex::new(r"(?i)publishdate|pubdate|timestamp|article_date|articledate|date").unwrap();
    }

    let mut date: Option<String> = None;

    'initial: for time in html.find(Name("time")) {
        if let Some(dt) = time.attr("datetime") {
            { date = Some(dt.to_string());
              break 'initial; }
        } else if let Some("timestamp") = time.attr("class") {
            { date = Some(time.text().trim_matches('\n').to_string());
              break 'initial; }
        }
    }

    if date.is_none() {
        'outer: for tag in html.find(Name("span")) {
            if let Some("datePublished") = tag.attr("itemprop") {
                if let Some(v) = tag.attr("content") {
                    { date = Some(v.to_string());
                      break 'outer; }
                } else if tag.text() != "" {
                    { date = Some(tag.text().trim_matches('\n').to_string());
                      break 'outer; }
                }
            }
        }
    }

    // These next three loops are due to the lack of `find_all` method for select.rs library
    // See also: https://github.com/Webhose/article-date-extractor/blob/master/articleDateExtractor/__init__.py#L191
    if date.is_none() {
        'outer_first: for tag in html.find(Name("span")) {
            if TAG_RE.is_match(tag.attr("class").unwrap_or("")) {
                { date = Some(tag.text().trim_matches('\n').to_string());
                  break 'outer_first; }
            }
        }
    }

    if date.is_none() {
        'outer_second: for tag in html.find(Name("p")) {
            if TAG_RE.is_match(tag.attr("class").unwrap_or("")) {
                { date = Some(tag.text().trim_matches('\n').to_string());
                  break 'outer_second; }
            }
        }
    }

    if date.is_none() {
        'outer_third: for tag in html.find(Name("div")) {
            if TAG_RE.is_match(tag.attr("class").unwrap_or("")) {
                { date = Some(tag.text().trim_matches('\n').to_string());
                  break 'outer_third; }
            }
        }
    }

    date
}

// Try to extract the date by using each function one by one
pub fn extract_article_published_date(link: &str, html: Option<String>) -> Result<NaiveDate> {
    let mut body: String = String::new();
    let mut parsed_body: Option<Document> = None;

    if let Some(v) = extract_from_url(link) {
        return parse_date(v.as_str())
    }

    if html.is_none() {
        if let Ok(mut response) = reqwest::get(link) {
            response.read_to_string(&mut body).unwrap();
            let doc = Document::from(body.as_str());
            parsed_body = Some(doc);
        } else {
            return Err("Couldn't open the link".into())
        }
    } else {
        parsed_body = Some(Document::from(html.unwrap().as_str()))
    }

    if let Some(v) = extract_from_url(link) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_ldjson(parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_meta(parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_html_tag(parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else {
        return Err("Couldn't find the date to parse".into())
    }
}

// Unit tests
#[cfg(test)]
mod test {
    use super::extract_from_url;
    use super::parse_date;
    use super::extract_from_meta;
    use super::extract_from_ldjson;
    use super::extract_from_html_tag;
    use super::extract_article_published_date;
    use chrono::NaiveDate;
    use reqwest;
    use std::string::String;
    use std::io::Read;
    use select::document::Document;

    #[test]
    fn parsing_date() {
        assert_eq!(NaiveDate::from_ymd(2015,11,30), parse_date("/2015/11/30/").unwrap());
        assert_eq!(NaiveDate::from_ymd(2015,11,30), parse_date("/2015/30/11/").unwrap());

        assert!(parse_date("bad_format").is_err());
    }

    #[test]
    fn extracting_from_url() {
        let link = "http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html";
        assert_eq!(Some("/2015/11/28/".to_string()), extract_from_url(link));

        let link = "";
        assert_eq!(None, extract_from_url(link));
    }

    #[test]
    fn extracting_from_meta() {
        let mut response = reqwest::get("https://techcrunch.com/2015/11/30/atlassian-share-price/").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();
        let document = Document::from(body.as_str());

        assert_eq!(Some(("2015-11-30 23:50:48".to_string())), extract_from_meta(&document));
    }

    #[test]
    fn extracting_from_ldjson() {
        let mut response = reqwest::get("https://techcrunch.com/2015/11/30/atlassian-share-price/").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();
        let document = Document::from(body.as_str());

        assert_eq!(Some("2015-12-01T07:50:48Z".to_string()), extract_from_ldjson(&document));
    }

    #[test]
    fn extracting_from_html_tag() {
        let mut response = reqwest::get("http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();
        let document = Document::from(body.as_str());

        assert_eq!(Some("2015-11-29T00:44:59Z".to_string()), extract_from_html_tag(&document));
    }

    #[test]
    fn extracting_article_published_date() {
        let link = "http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html";
        let mut response = reqwest::get("http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();

        assert_eq!(NaiveDate::from_ymd(2015,11,28), extract_article_published_date(&link, None).unwrap());
        assert_eq!(NaiveDate::from_ymd(2015,11,28), extract_article_published_date(&link, Some(body)).unwrap());

        assert!((extract_article_published_date("", None)).is_err());
    }
}
<|endoftext|>


<commit_before>use std::rc::Rc;
use std::cell::RefCell;

use progress::{Timestamp, Operate, Antichain};
use progress::frontier::MutableAntichain;
use progress::nested::subgraph::Source::ChildOutput;
use progress::nested::subgraph::Target::ChildInput;
use progress::count_map::CountMap;

use dataflow::channels::pushers::Tee;
use dataflow::channels::pushers::Counter as PushCounter;
use dataflow::channels::pushers::buffer::Buffer as PushBuffer;
use dataflow::channels::pact::{ParallelizationContract, Pipeline};
use dataflow::channels::pullers::Counter as PullCounter;

use Data;
use dataflow::{Stream, Scope};

pub trait Probe<G: Scope, D: Data> {
    /// Constructs a progress probe which can indicates which timestamps have elapsed at the operator.
    fn probe(&self) -> (Handle<G::Timestamp>, Stream<G, D>);
}

impl<G: Scope, D: Data> Probe<G, D> for Stream<G, D> {
    fn probe(&self) -> (Handle<G::Timestamp>, Stream<G, D>) {

        // the frontier is shared state; scope updates, handle reads.
        let frontier = Rc::new(RefCell::new(MutableAntichain::new()));
        //
        // // we add the scope, acquiring the name of the probe, then add an edge.
        // let index = self.scope().add_operator(Operator { frontier: frontier.clone() });
        // self.scope().add_edge(*self.name(), ChildInput(index, 0));

        // the handle is the only result
        let handle = Handle { frontier: frontier.clone() };

        let mut scope = self.scope();   // clones the scope

        let (sender, receiver) = Pipeline.connect(&mut scope);
        let (targets, registrar) = Tee::<G::Timestamp,D>::new();
        let operator = Operator {
            input: PullCounter::new(receiver),
            output: PushBuffer::new(PushCounter::new(targets, Rc::new(RefCell::new(CountMap::new())))),
            frontier: frontier,
        };

        let index = scope.add_operator(operator);
        self.connect_to(ChildInput(index, 0), sender);
        (handle, Stream::new(ChildOutput(index, 0), registrar, scope))
    }
}

#[derive(Default)]
pub struct Handle<T:Timestamp> {
    frontier: Rc<RefCell<MutableAntichain<T>>>
}

impl<T: Timestamp> Handle<T> {
    #[inline] pub fn lt(&self, time: &T) -> bool { self.frontier.borrow().lt(time) }
    #[inline] pub fn le(&self, time: &T) -> bool { self.frontier.borrow().le(time) }
    #[inline] pub fn done(&self) -> bool { self.frontier.borrow().elements().len() == 0 }
}

struct Operator<T:Timestamp, D: Data> {
    input: PullCounter<T, D>,
    output: PushBuffer<T, D, PushCounter<T, D, Tee<T, D>>>,
    frontier: Rc<RefCell<MutableAntichain<T>>>
}

impl<T:Timestamp, D: Data> Operate<T> for Operator<T, D> {
    fn name(&self) -> &str { "Probe" }
    fn inputs(&self) -> usize { 1 }
    fn outputs(&self) -> usize { 1 }

    // we need to set the initial value of the frontier
    fn set_external_summary(&mut self, _: Vec<Vec<Antichain<T::Summary>>>, counts: &mut [CountMap<T>]) {
        let mut borrow = self.frontier.borrow_mut();
        while let Some((time, delta)) = counts[0].pop() {
            borrow.update(&time, delta);
        }
    }

    // each change to the frontier should be shared
    fn push_external_progress(&mut self, counts: &mut [CountMap<T>]) {
        let mut borrow = self.frontier.borrow_mut();
        while let Some((time, delta)) = counts[0].pop() {
            borrow.update(&time, delta);
        }
    }

    // the scope does nothing. this is actually a problem, because "reachability" assumes all messages on each edge.
    fn pull_internal_progress(&mut self,_: &mut [CountMap<T>], consumed: &mut [CountMap<T>], produced: &mut [CountMap<T>]) -> bool {

        while let Some((time, data)) = self.input.next() {
            self.output.session(time).give_content(data);
        }

        self.output.cease();

        // extract what we know about progress from the input and output adapters.
        self.input.pull_progress(&mut consumed[0]);
        self.output.inner().pull_progress(&mut produced[0]);

        false
    }

    fn notify_me(&self) -> bool { true }
}

#[cfg(test)]
mod tests {

    use ::Configuration;
    use ::progress::timestamp::RootTimestamp;
    use dataflow::*;
    use dataflow::operators::{Input, Probe};

    #[test]
    fn probe() {

        // initializes and runs a timely dataflow computation
        ::execute(Configuration::Thread, |computation| {

            // create a new input, and inspect its output
            let (mut input, probe) = computation.scoped(move |builder| {
                let (input, stream) = builder.new_input::<String>();
                (input, stream.probe().0)
            });

            // introduce data and watch!
            for round in 0..10 {
                assert!(!probe.done());
                assert!(probe.le(&RootTimestamp::new(round)));
                assert!(!probe.lt(&RootTimestamp::new(round)));
                assert!(probe.lt(&RootTimestamp::new(round + 1)));
                input.advance_to(round + 1);
                computation.step();
            }

            // seal the input
            input.close();

            // finish off any remaining work
            computation.step();
            assert!(probe.done());
        });
    }

}
<commit_msg>documentation<commit_after>//! An extension method to monitor progress at a `Stream`.

use std::rc::Rc;
use std::cell::RefCell;

use progress::{Timestamp, Operate, Antichain};
use progress::frontier::MutableAntichain;
use progress::nested::subgraph::Source::ChildOutput;
use progress::nested::subgraph::Target::ChildInput;
use progress::count_map::CountMap;

use dataflow::channels::pushers::Tee;
use dataflow::channels::pushers::Counter as PushCounter;
use dataflow::channels::pushers::buffer::Buffer as PushBuffer;
use dataflow::channels::pact::{ParallelizationContract, Pipeline};
use dataflow::channels::pullers::Counter as PullCounter;

use Data;
use dataflow::{Stream, Scope};

/// Monitors progress at a `Stream`.
pub trait Probe<G: Scope, D: Data> {
    /// Constructs a progress probe which indicates which timestamps have elapsed at the operator.
    fn probe(&self) -> (Handle<G::Timestamp>, Stream<G, D>);
}

impl<G: Scope, D: Data> Probe<G, D> for Stream<G, D> {
    fn probe(&self) -> (Handle<G::Timestamp>, Stream<G, D>) {

        // the frontier is shared state; scope updates, handle reads.
        let frontier = Rc::new(RefCell::new(MutableAntichain::new()));
        let handle = Handle { frontier: frontier.clone() };

        let mut scope = self.scope();   // clones the scope

        let (sender, receiver) = Pipeline.connect(&mut scope);
        let (targets, registrar) = Tee::<G::Timestamp,D>::new();
        let operator = Operator {
            input: PullCounter::new(receiver),
            output: PushBuffer::new(PushCounter::new(targets, Rc::new(RefCell::new(CountMap::new())))),
            frontier: frontier,
        };

        let index = scope.add_operator(operator);
        self.connect_to(ChildInput(index, 0), sender);
        (handle, Stream::new(ChildOutput(index, 0), registrar, scope))
    }
}

/// Reports information about progress at the probe.
pub struct Handle<T:Timestamp> {
    frontier: Rc<RefCell<MutableAntichain<T>>>
}

impl<T: Timestamp> Handle<T> {
    /// returns true iff the frontier is strictly less than `time`.
    #[inline] pub fn lt(&self, time: &T) -> bool { self.frontier.borrow().lt(time) }
    /// returns true iff the frontier is less than or equal to `time`.
    #[inline] pub fn le(&self, time: &T) -> bool { self.frontier.borrow().le(time) }
    /// returns true iff the frontier is empty.
    #[inline] pub fn done(&self) -> bool { self.frontier.borrow().elements().len() == 0 }
}

struct Operator<T:Timestamp, D: Data> {
    input: PullCounter<T, D>,
    output: PushBuffer<T, D, PushCounter<T, D, Tee<T, D>>>,
    frontier: Rc<RefCell<MutableAntichain<T>>>
}

impl<T:Timestamp, D: Data> Operate<T> for Operator<T, D> {
    fn name(&self) -> &str { "Probe" }
    fn inputs(&self) -> usize { 1 }
    fn outputs(&self) -> usize { 1 }

    // we need to set the initial value of the frontier
    fn set_external_summary(&mut self, _: Vec<Vec<Antichain<T::Summary>>>, counts: &mut [CountMap<T>]) {
        let mut borrow = self.frontier.borrow_mut();
        while let Some((time, delta)) = counts[0].pop() {
            borrow.update(&time, delta);
        }
    }

    // each change to the frontier should be shared
    fn push_external_progress(&mut self, counts: &mut [CountMap<T>]) {
        let mut borrow = self.frontier.borrow_mut();
        while let Some((time, delta)) = counts[0].pop() {
            borrow.update(&time, delta);
        }
    }

    // the scope does nothing. this is actually a problem, because "reachability" assumes all messages on each edge.
    fn pull_internal_progress(&mut self,_: &mut [CountMap<T>], consumed: &mut [CountMap<T>], produced: &mut [CountMap<T>]) -> bool {

        while let Some((time, data)) = self.input.next() {
            self.output.session(time).give_content(data);
        }

        self.output.cease();

        // extract what we know about progress from the input and output adapters.
        self.input.pull_progress(&mut consumed[0]);
        self.output.inner().pull_progress(&mut produced[0]);

        false
    }

    fn notify_me(&self) -> bool { true }
}

#[cfg(test)]
mod tests {

    use ::Configuration;
    use ::progress::timestamp::RootTimestamp;
    use dataflow::*;
    use dataflow::operators::{Input, Probe};

    #[test]
    fn probe() {

        // initializes and runs a timely dataflow computation
        ::execute(Configuration::Thread, |computation| {

            // create a new input, and inspect its output
            let (mut input, probe) = computation.scoped(move |builder| {
                let (input, stream) = builder.new_input::<String>();
                (input, stream.probe().0)
            });

            // introduce data and watch!
            for round in 0..10 {
                assert!(!probe.done());
                assert!(probe.le(&RootTimestamp::new(round)));
                assert!(!probe.lt(&RootTimestamp::new(round)));
                assert!(probe.lt(&RootTimestamp::new(round + 1)));
                input.advance_to(round + 1);
                computation.step();
            }

            // seal the input
            input.close();

            // finish off any remaining work
            computation.step();
            assert!(probe.done());
        });
    }

}
<|endoftext|>


<commit_before>use distribution;
use random;

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    k: usize,
    p: Vec<f64>,
    cumsum: Vec<f64>,
}

impl Categorical {
    /// Create a categorical distribution with success probability `p`.
    ///
    /// It should hold that `p[i] >= 0`, `p[i] <= 1`, and `sum(p) == 1`.
    pub fn new(p: &[f64]) -> Categorical {
        should!(is_probability_vector(p), {
            const EPSILON: f64 = 1e-12;
            p.iter().all(|&p| p >= 0.0 && p <= 1.0) &&
                (p.iter().fold(0.0, |sum, &p| sum + p) - 1.0).abs() < EPSILON
        });

        let k = p.len();
        let mut cumsum = p.to_vec();
        for i in 1..(k - 1) {
            cumsum[i] += cumsum[i - 1];
        }
        cumsum[k - 1] = 1.0;
        Categorical { k: k, p: p.to_vec(), cumsum: cumsum }
    }

    /// Return the number of categories.
    #[inline(always)]
    pub fn k(&self) -> usize { self.k }

    /// Return the event probabilities.
    #[inline(always)]
    pub fn p(&self) -> &[f64] { &self.p }
}

impl distribution::Distribution for Categorical {
    type Value = usize;

    fn cdf(&self, x: f64) -> f64 {
        if x < 0.0 {
            return 0.0;
        }
        let x = x as usize;
        if x >= self.k {
            return 1.0;
        }
        self.cumsum[x]
    }
}

impl distribution::Discrete for Categorical {
    #[inline]
    fn pmf(&self, x: usize) -> f64 {
        should!(x < self.k);
        self.p[x]
    }
}

impl distribution::Entropy for Categorical {
    fn entropy(&self) -> f64 {
        -self.p.iter().fold(0.0, |sum, p| sum + p * p.ln())
    }
}

impl distribution::Inverse for Categorical {
    fn inv_cdf(&self, p: f64) -> usize {
        should!(0.0 <= p && p <= 1.0);
        self.cumsum.iter().position(|&sum| sum > 0.0 && sum >= p).unwrap_or_else(|| {
            self.p.iter().rposition(|&p| p > 0.0).unwrap()
        })
    }
}

impl distribution::Kurtosis for Categorical {
    fn kurtosis(&self) -> f64 {
        use distribution::{Mean, Variance};
        let (mean, variance) = (self.mean(), self.variance());
        let kurt = self.p.iter().enumerate().fold(0.0, |sum, (i, p)| {
            sum + (i as f64 - mean).powi(4) * p
        });
        kurt / variance.powi(2) - 3.0
    }
}

impl distribution::Mean for Categorical {
    fn mean(&self) -> f64 {
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + i as f64 * p)
    }
}

impl distribution::Median for Categorical {
    fn median(&self) -> f64 {
        if self.p[0] > 0.5 {
            return 0.0;
        }
        if self.p[0] == 0.5 {
            return 0.5;
        }
        for (i, &sum) in self.cumsum.iter().enumerate() {
            if sum == 0.5 {
                return (2 * i - 1) as f64 / 2.0;
            } else if sum > 0.5 {
                return i as f64;
            }
        }
        unreachable!()
    }
}

impl distribution::Modes for Categorical {
    fn modes(&self) -> Vec<usize> {
        let mut modes = Vec::new();
        let mut max = 0.0;
        for (i, &p) in self.p.iter().enumerate() {
            if p == max {
                modes.push(i);
            }
            if p > max {
                max = p;
                modes = vec![i];
            }
        }
        modes
    }
}

impl distribution::Sample for Categorical {
    #[inline]
    fn sample<S>(&self, source: &mut S) -> usize where S: random::Source {
        use distribution::Inverse;
        self.inv_cdf(source.read::<f64>())
    }
}

impl distribution::Skewness for Categorical {
    fn skewness(&self) -> f64 {
        use distribution::{Mean, Variance};
        let (mean, variance) = (self.mean(), self.variance());
        let skew = self.p.iter().enumerate().fold(0.0, |sum, (i, p)| {
            sum + (i as f64 - mean).powi(3) * p
        });
        skew / (variance * variance.sqrt())
    }
}

impl distribution::Variance for Categorical {
    fn variance(&self) -> f64 {
        use distribution::Mean;
        let mean = self.mean();
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| {
            sum + (i as f64 - mean).powi(2) * p
        })
    }
}

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

    macro_rules! new(
        (equal $k:expr) => { Categorical::new(&[1.0 / $k as f64; $k]) };
        ($p:expr) => { Categorical::new(&$p); }
    );

    #[test]
    fn cdf() {
        let d = new!([0.0, 0.75, 0.25, 0.0]);
        let p = vec![0.0, 0.0, 0.75, 1.0, 1.0];

        let x = (-1..4).map(|x| d.cdf(x as f64)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let x = (-1..4).map(|x| d.cdf(x as f64 + 0.5)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let d = new!(equal 3);
        let p = vec![0.0, 1.0 / 3.0, 2.0 / 3.0, 1.0];

        let x = (-1..3).map(|x| d.cdf(x as f64)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let x = (-1..3).map(|x| d.cdf(x as f64 + 0.5)).collect::<Vec<_>>();
        assert_eq!(&x, &p);
    }

    #[test]
    fn pmf() {
        let p = [0.0, 0.75, 0.25, 0.0];
        let d = new!(p);
        assert_eq!(&(0..4).map(|x| d.pmf(x)).collect::<Vec<_>>(), &p.to_vec());

        let d = new!(equal 3);
        assert_eq!(&(0..3).map(|x| d.pmf(x)).collect::<Vec<_>>(), &vec![1.0 / 3.0; 3])
    }

    #[test]
    fn entropy() {
        use std::f64::consts::LN_2;
        assert_eq!(new!(equal 2).entropy(), LN_2);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).entropy(), 1.2798542258336676);
    }

    #[test]
    fn inv_cdf() {
        let d = new!([0.0, 0.75, 0.25, 0.0]);
        let p = vec![0.0, 0.75, 0.7500001, 1.0];
        assert_eq!(&p.iter().map(|&p| d.inv_cdf(p)).collect::<Vec<_>>(), &vec![1, 1, 2, 2]);

        let d = new!(equal 3);
        let p = vec![0.0, 0.5, 0.75, 1.0];
        assert_eq!(&p.iter().map(|&p| d.inv_cdf(p)).collect::<Vec<_>>(), &vec![0, 1, 2, 2]);
    }

    #[test]
    fn kurtosis() {
        assert_eq!(new!(equal 2).kurtosis(), -2.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).kurtosis(), -0.7999999999999998);
    }

    #[test]
    fn mean() {
        assert_eq!(new!(equal 3).mean(), 1.0);
        assert_eq!(new!([0.3, 0.3, 0.4]).mean(), 1.1);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).mean(), 1.5);
    }

    #[test]
    fn median() {
        assert_eq!(new!([0.6, 0.2, 0.2]).median(), 0.0);
        assert_eq!(new!(equal 2).median(), 0.5);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).median(), 2.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).median(), 0.5);
    }

    #[test]
    fn modes() {
        assert_eq!(new!([0.6, 0.2, 0.2]).modes(), vec![0]);
        assert_eq!(new!(equal 2).modes(), vec![0, 1]);
        assert_eq!(new!(equal 3).modes(), vec![0, 1, 2]);
        assert_eq!(new!([0.4, 0.2, 0.4]).modes(), vec![0, 2]);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).modes(), vec![1, 2]);
    }

    #[test]
    fn sample() {
        let mut source = random::default();
        let sum = Independent(&new!([0.0, 0.5, 0.5]), &mut source).take(100).fold(0, |a, b| a + b);
        assert!(100 <= sum && sum <= 200);
    }

    #[test]
    fn skewness() {
        assert_eq!(new!(equal 6).skewness(), 0.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).skewness(), 0.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).skewness(), -0.6);
    }

    #[test]
    fn variance() {
        assert_eq!(new!(equal 3).variance(), 2.0 / 3.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).variance(), 11.0 / 12.0);
    }
}
<commit_msg>categorical: extend the sample test<commit_after>use distribution;
use random;

/// A categorical distribution.
#[derive(Clone)]
pub struct Categorical {
    k: usize,
    p: Vec<f64>,
    cumsum: Vec<f64>,
}

impl Categorical {
    /// Create a categorical distribution with success probability `p`.
    ///
    /// It should hold that `p[i] >= 0`, `p[i] <= 1`, and `sum(p) == 1`.
    pub fn new(p: &[f64]) -> Categorical {
        should!(is_probability_vector(p), {
            const EPSILON: f64 = 1e-12;
            p.iter().all(|&p| p >= 0.0 && p <= 1.0) &&
                (p.iter().fold(0.0, |sum, &p| sum + p) - 1.0).abs() < EPSILON
        });

        let k = p.len();
        let mut cumsum = p.to_vec();
        for i in 1..(k - 1) {
            cumsum[i] += cumsum[i - 1];
        }
        cumsum[k - 1] = 1.0;
        Categorical { k: k, p: p.to_vec(), cumsum: cumsum }
    }

    /// Return the number of categories.
    #[inline(always)]
    pub fn k(&self) -> usize { self.k }

    /// Return the event probabilities.
    #[inline(always)]
    pub fn p(&self) -> &[f64] { &self.p }
}

impl distribution::Distribution for Categorical {
    type Value = usize;

    fn cdf(&self, x: f64) -> f64 {
        if x < 0.0 {
            return 0.0;
        }
        let x = x as usize;
        if x >= self.k {
            return 1.0;
        }
        self.cumsum[x]
    }
}

impl distribution::Discrete for Categorical {
    #[inline]
    fn pmf(&self, x: usize) -> f64 {
        should!(x < self.k);
        self.p[x]
    }
}

impl distribution::Entropy for Categorical {
    fn entropy(&self) -> f64 {
        -self.p.iter().fold(0.0, |sum, p| sum + p * p.ln())
    }
}

impl distribution::Inverse for Categorical {
    fn inv_cdf(&self, p: f64) -> usize {
        should!(0.0 <= p && p <= 1.0);
        self.cumsum.iter().position(|&sum| sum > 0.0 && sum >= p).unwrap_or_else(|| {
            self.p.iter().rposition(|&p| p > 0.0).unwrap()
        })
    }
}

impl distribution::Kurtosis for Categorical {
    fn kurtosis(&self) -> f64 {
        use distribution::{Mean, Variance};
        let (mean, variance) = (self.mean(), self.variance());
        let kurt = self.p.iter().enumerate().fold(0.0, |sum, (i, p)| {
            sum + (i as f64 - mean).powi(4) * p
        });
        kurt / variance.powi(2) - 3.0
    }
}

impl distribution::Mean for Categorical {
    fn mean(&self) -> f64 {
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| sum + i as f64 * p)
    }
}

impl distribution::Median for Categorical {
    fn median(&self) -> f64 {
        if self.p[0] > 0.5 {
            return 0.0;
        }
        if self.p[0] == 0.5 {
            return 0.5;
        }
        for (i, &sum) in self.cumsum.iter().enumerate() {
            if sum == 0.5 {
                return (2 * i - 1) as f64 / 2.0;
            } else if sum > 0.5 {
                return i as f64;
            }
        }
        unreachable!()
    }
}

impl distribution::Modes for Categorical {
    fn modes(&self) -> Vec<usize> {
        let mut modes = Vec::new();
        let mut max = 0.0;
        for (i, &p) in self.p.iter().enumerate() {
            if p == max {
                modes.push(i);
            }
            if p > max {
                max = p;
                modes = vec![i];
            }
        }
        modes
    }
}

impl distribution::Sample for Categorical {
    #[inline]
    fn sample<S>(&self, source: &mut S) -> usize where S: random::Source {
        use distribution::Inverse;
        self.inv_cdf(source.read::<f64>())
    }
}

impl distribution::Skewness for Categorical {
    fn skewness(&self) -> f64 {
        use distribution::{Mean, Variance};
        let (mean, variance) = (self.mean(), self.variance());
        let skew = self.p.iter().enumerate().fold(0.0, |sum, (i, p)| {
            sum + (i as f64 - mean).powi(3) * p
        });
        skew / (variance * variance.sqrt())
    }
}

impl distribution::Variance for Categorical {
    fn variance(&self) -> f64 {
        use distribution::Mean;
        let mean = self.mean();
        self.p.iter().enumerate().fold(0.0, |sum, (i, p)| {
            sum + (i as f64 - mean).powi(2) * p
        })
    }
}

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

    macro_rules! new(
        (equal $k:expr) => { Categorical::new(&[1.0 / $k as f64; $k]) };
        ($p:expr) => { Categorical::new(&$p); }
    );

    #[test]
    fn cdf() {
        let d = new!([0.0, 0.75, 0.25, 0.0]);
        let p = vec![0.0, 0.0, 0.75, 1.0, 1.0];

        let x = (-1..4).map(|x| d.cdf(x as f64)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let x = (-1..4).map(|x| d.cdf(x as f64 + 0.5)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let d = new!(equal 3);
        let p = vec![0.0, 1.0 / 3.0, 2.0 / 3.0, 1.0];

        let x = (-1..3).map(|x| d.cdf(x as f64)).collect::<Vec<_>>();
        assert_eq!(&x, &p);

        let x = (-1..3).map(|x| d.cdf(x as f64 + 0.5)).collect::<Vec<_>>();
        assert_eq!(&x, &p);
    }

    #[test]
    fn pmf() {
        let p = [0.0, 0.75, 0.25, 0.0];
        let d = new!(p);
        assert_eq!(&(0..4).map(|x| d.pmf(x)).collect::<Vec<_>>(), &p.to_vec());

        let d = new!(equal 3);
        assert_eq!(&(0..3).map(|x| d.pmf(x)).collect::<Vec<_>>(), &vec![1.0 / 3.0; 3])
    }

    #[test]
    fn entropy() {
        use std::f64::consts::LN_2;
        assert_eq!(new!(equal 2).entropy(), LN_2);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).entropy(), 1.2798542258336676);
    }

    #[test]
    fn inv_cdf() {
        let d = new!([0.0, 0.75, 0.25, 0.0]);
        let p = vec![0.0, 0.75, 0.7500001, 1.0];
        assert_eq!(&p.iter().map(|&p| d.inv_cdf(p)).collect::<Vec<_>>(), &vec![1, 1, 2, 2]);

        let d = new!(equal 3);
        let p = vec![0.0, 0.5, 0.75, 1.0];
        assert_eq!(&p.iter().map(|&p| d.inv_cdf(p)).collect::<Vec<_>>(), &vec![0, 1, 2, 2]);
    }

    #[test]
    fn kurtosis() {
        assert_eq!(new!(equal 2).kurtosis(), -2.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).kurtosis(), -0.7999999999999998);
    }

    #[test]
    fn mean() {
        assert_eq!(new!(equal 3).mean(), 1.0);
        assert_eq!(new!([0.3, 0.3, 0.4]).mean(), 1.1);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).mean(), 1.5);
    }

    #[test]
    fn median() {
        assert_eq!(new!([0.6, 0.2, 0.2]).median(), 0.0);
        assert_eq!(new!(equal 2).median(), 0.5);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).median(), 2.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).median(), 0.5);
    }

    #[test]
    fn modes() {
        assert_eq!(new!([0.6, 0.2, 0.2]).modes(), vec![0]);
        assert_eq!(new!(equal 2).modes(), vec![0, 1]);
        assert_eq!(new!(equal 3).modes(), vec![0, 1, 2]);
        assert_eq!(new!([0.4, 0.2, 0.4]).modes(), vec![0, 2]);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).modes(), vec![1, 2]);
    }

    #[test]
    fn sample() {
        let mut source = random::default();

        let sum = Independent(&new!([0.0, 0.5, 0.5]), &mut source).take(100).fold(0, |a, b| a + b);
        assert!(100 <= sum && sum <= 200);

        let p = (0..11).map(|i| if i % 2 != 0 { 0.2 } else { 0.0 }).collect::<Vec<_>>();
        assert!(Independent(&new!(p), &mut source).take(1000).all(|x| x % 2 != 0));
    }

    #[test]
    fn skewness() {
        assert_eq!(new!(equal 6).skewness(), 0.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).skewness(), 0.0);
        assert_eq!(new!([0.1, 0.2, 0.3, 0.4]).skewness(), -0.6);
    }

    #[test]
    fn variance() {
        assert_eq!(new!(equal 3).variance(), 2.0 / 3.0);
        assert_eq!(new!([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]).variance(), 11.0 / 12.0);
    }
}
<|endoftext|>