<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

extern crate curl;
extern crate rs_es;
extern crate serde;
extern crate serde_json;
extern crate regex;

use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};

use super::objects::{Addr, Incr, DocType};

// Rubber is an wrapper around elasticsearch API
pub struct Rubber {
    index_name: String,
    client: rs_es::Client,
}

impl Rubber {
    // build a rubber with a connection string (http://host:port/index)
    pub fn new(cnx: &str) -> Rubber {
        let re = regex::Regex::new(r"(?:https?://)?(?P<host>.+?):(?P<port>\d+)/(?P<index>\w+)")
                     .unwrap();
        let cap = re.captures(cnx).unwrap();
        let host = cap.name("host").unwrap();
        let port = cap.name("port").unwrap().parse::<u32>().unwrap();
        let index = cap.name("index").unwrap();
        info!("elastic search host {:?} port {:?} index {:?}",
              host,
              port,
              index);

        Rubber {
            index_name: index.to_string(),
            client: rs_es::Client::new(&host, port),
        }
    }
    pub fn clean_db_by_doc_type(&mut self,
                                doc_type: &[&str])
                                -> Result<usize, rs_es::error::EsError> {
        info!("Clean elasticsearch db...");
        let scroll = rs_es::units::Duration::minutes(1);
        let mut scan: rs_es::operations::search::ScanResult<serde_json::Value> =
            match self.client
                      .search_query()
                      .with_indexes(&[&self.index_name])
                      .with_size(10000)
                      .with_source(rs_es::operations::search::Source::include(&["_id"]))
                      .with_types(&doc_type)
                      .scan(&scroll) {
                Ok(scan) => scan,
                Err(e) => {
                    info!("Scan error: {:?}", e);
                    return Err(e);
                }
            };
        let mut count: usize = 0;
        loop {
            let page = match scan.scroll(&mut self.client, &scroll) {
                Ok(page) => page,
                Err(e) => {
                    info!("scroll error: {:?}", e);
                    try!(scan.close(&mut self.client));
                    return Err(e);
                }
            };
            let mut hits = page.hits.hits;
            if hits.len() == 0 {
                break;
            }
            count = count + hits.len();
            info!("Delete : Count {}", hits.len());
            let actions: Vec<rs_es::operations::bulk::Action<()>> =
                hits.drain(..)
                    .map(|hit| {
                        rs_es::operations::bulk::Action::delete(hit.id)
                            .with_index(hit.index)
                            .with_doc_type(hit.doc_type)
                    })
                    .collect();
            try!(self.client.bulk(&actions).send());
        }
        try!(scan.close(&mut self.client));
        Ok((count))
    }

    pub fn create_index(&mut self) {
        debug!("creating index");
        // Note: for the moment I don't see an easy way to do this with rs_es
        let analysis = include_str!("../json/settings.json");
        let res = curl::http::handle()
                      .put(self.client.full_url(&self.index_name), analysis)
                      .exec()
                      .map(|res| res.get_code() == 200)
                      .unwrap_or(false);
        if !res {
            info!("Error adding analysis");
        }
    }

    pub fn delete_index(&mut self) -> Result<(), Box<::std::error::Error>> {
        debug!("deleting index");
        let res = self.client
                      .delete_index(&self.index_name)
                      .map(|res| res.acknowledged)
                      .unwrap_or(false);
        if !res {
            Err(format!("Error deleting index {}", &self.index_name).into())
        } else {
            Ok(())
        }
    }

    pub fn bulk_index<T, I>(&mut self, mut iter: I) -> Result<u32, rs_es::error::EsError>
        where T: serde::Serialize + DocType,
              I: Iterator<Item = T>
    {
        use self::rs_es::operations::bulk::Action;
        let mut chunk = Vec::new();
        let mut nb = 0;
        loop {
            chunk.clear();
            let addr = match iter.next() {
                Some(a) => a,
                None => break,
            };
            chunk.push(Action::index(addr));

            nb += 1;
            for addr in iter.by_ref().take(1000) {
                chunk.push(Action::index(addr));
                nb += 1;
            }
            try!(self.client
                     .bulk(&chunk)
                     .with_index(&self.index_name)
                     .with_doc_type(T::doc_type())
                     .send());
        }

        Ok(nb)
    }

    pub fn index<I: Iterator<Item = Addr>>(&mut self,
                                           iter: I)
                                           -> Result<u32, rs_es::error::EsError> {
        let mut admins = HashMap::new();
        let mut streets = HashMap::new();
        let mut nb = 0;

        nb += try!(self.bulk_index(iter.inspect(|addr| {
            upsert(&addr.street.administrative_region, &mut admins);
            upsert(&addr.street, &mut streets);
        })));
        nb += try!(self.bulk_index(admins.values()));
        nb += try!(self.bulk_index(streets.values()));
        Ok(nb)
    }
}

fn upsert<T: Incr>(elt: &T, map: &mut HashMap<String, T>) {
    match map.entry(elt.id().to_string()) {
        Vacant(e) => {
            e.insert(elt.clone());
        }
        Occupied(mut e) => e.get_mut().incr(),
    }
}
<commit_msg>Don't import street and admin from bano, we want to take them from osm<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

extern crate curl;
extern crate rs_es;
extern crate serde;
extern crate serde_json;
extern crate regex;

use super::objects::{Addr, Incr, DocType};

// Rubber is an wrapper around elasticsearch API
pub struct Rubber {
    index_name: String,
    client: rs_es::Client,
}

impl Rubber {
    // build a rubber with a connection string (http://host:port/index)
    pub fn new(cnx: &str) -> Rubber {
        let re = regex::Regex::new(r"(?:https?://)?(?P<host>.+?):(?P<port>\d+)/(?P<index>\w+)")
                     .unwrap();
        let cap = re.captures(cnx).unwrap();
        let host = cap.name("host").unwrap();
        let port = cap.name("port").unwrap().parse::<u32>().unwrap();
        let index = cap.name("index").unwrap();
        info!("elastic search host {:?} port {:?} index {:?}",
              host,
              port,
              index);

        Rubber {
            index_name: index.to_string(),
            client: rs_es::Client::new(&host, port),
        }
    }
    pub fn clean_db_by_doc_type(&mut self,
                                doc_type: &[&str])
                                -> Result<usize, rs_es::error::EsError> {
        info!("Clean elasticsearch db...");
        let scroll = rs_es::units::Duration::minutes(1);
        let mut scan: rs_es::operations::search::ScanResult<serde_json::Value> =
            match self.client
                      .search_query()
                      .with_indexes(&[&self.index_name])
                      .with_size(10000)
                      .with_source(rs_es::operations::search::Source::include(&["_id"]))
                      .with_types(&doc_type)
                      .scan(&scroll) {
                Ok(scan) => scan,
                Err(e) => {
                    info!("Scan error: {:?}", e);
                    return Err(e);
                }
            };
        let mut count: usize = 0;
        loop {
            let page = match scan.scroll(&mut self.client, &scroll) {
                Ok(page) => page,
                Err(e) => {
                    info!("scroll error: {:?}", e);
                    try!(scan.close(&mut self.client));
                    return Err(e);
                }
            };
            let mut hits = page.hits.hits;
            if hits.len() == 0 {
                break;
            }
            count = count + hits.len();
            info!("Delete : Count {}", hits.len());
            let actions: Vec<rs_es::operations::bulk::Action<()>> =
                hits.drain(..)
                    .map(|hit| {
                        rs_es::operations::bulk::Action::delete(hit.id)
                            .with_index(hit.index)
                            .with_doc_type(hit.doc_type)
                    })
                    .collect();
            try!(self.client.bulk(&actions).send());
        }
        try!(scan.close(&mut self.client));
        Ok((count))
    }

    pub fn create_index(&mut self) {
        debug!("creating index");
        // Note: for the moment I don't see an easy way to do this with rs_es
        let analysis = include_str!("../json/settings.json");
        let res = curl::http::handle()
                      .put(self.client.full_url(&self.index_name), analysis)
                      .exec()
                      .map(|res| res.get_code() == 200)
                      .unwrap_or(false);
        if !res {
            info!("Error adding analysis");
        }
    }

    pub fn delete_index(&mut self) -> Result<(), Box<::std::error::Error>> {
        debug!("deleting index");
        let res = self.client
                      .delete_index(&self.index_name)
                      .map(|res| res.acknowledged)
                      .unwrap_or(false);
        if !res {
            Err(format!("Error deleting index {}", &self.index_name).into())
        } else {
            Ok(())
        }
    }

    pub fn bulk_index<T, I>(&mut self, mut iter: I) -> Result<u32, rs_es::error::EsError>
        where T: serde::Serialize + DocType,
              I: Iterator<Item = T>
    {
        use self::rs_es::operations::bulk::Action;
        let mut chunk = Vec::new();
        let mut nb = 0;
        loop {
            chunk.clear();
            let addr = match iter.next() {
                Some(a) => a,
                None => break,
            };
            chunk.push(Action::index(addr));

            nb += 1;
            for addr in iter.by_ref().take(1000) {
                chunk.push(Action::index(addr));
                nb += 1;
            }
            try!(self.client
                     .bulk(&chunk)
                     .with_index(&self.index_name)
                     .with_doc_type(T::doc_type())
                     .send());
        }

        Ok(nb)
    }

    pub fn index<I: Iterator<Item = Addr>>(&mut self,
                                           iter: I)
                                           -> Result<u32, rs_es::error::EsError> {
        let mut nb = 0;
        nb += try!(self.bulk_index(iter));
        Ok(nb)
    }
}
<|endoftext|>


<commit_before>//! Run the static site generator.

mod cli;

// Standard library
use std::io::prelude::*;

// First party
use crate::cli::Command;

fn main() {
    if let Err(reason) = run() {
        // if this fails, we literally can't do a thing except panic.
        write!(std::io::stderr(), "failure: {}", reason).unwrap();
        std::process::exit(1);
    }
}

/// Define a `Result`-returning function to run the app.
///
/// (This is a standard Rust pattern to support the use of `try~`/`?`. We're
/// not doing that yet, but I expect we might eventually; this is convenient.)
fn run() -> Result<(), String> {
    let cwd = std::env::current_dir()
        .expect("Something is suuuuper borked: I cannot even get the current working directory!");

    match Command::cli() {
        Command::Build { site_directory } => lightning::build(site_directory.unwrap_or(cwd)),
    }
}
<commit_msg>Hey, `main` can use `Result` since forever now!<commit_after>//! Run the static site generator.

mod cli;

use crate::cli::Command;

fn main() -> Result<(), String> {
    let cwd = std::env::current_dir()
        .expect("Something is suuuuper borked: I cannot even get the current working directory!");

    match Command::cli() {
        Command::Build { site_directory } => lightning::build(site_directory.unwrap_or(cwd)),
    }
}
<|endoftext|>


<commit_before>// ---------------------------------------------------------------------------------------
// Ruthenium, an ack-like searcher, (c) 2015 Georg Brandl.
// Licensed under the MIT license.
// ---------------------------------------------------------------------------------------

use std::cmp::min;
use std::path::Path;

#[cfg(feature = "pcre")]
use pcre::Regex;
#[cfg(not(feature = "pcre"))]
use regex::bytes::Regex;

use options::{Casing, Opts};

/// Represents a line that matched the pattern (maybe multiple times).
#[derive(Debug)]
pub struct Match {
    /// Line number in the file
    pub lineno: usize,
    /// Line text
    pub line: Vec<u8>,
    /// Spans (start, end) of matching parts in the line
    pub spans: Vec<(usize, usize)>,
    /// Context lines before the matched line
    pub before: Vec<Vec<u8>>,
    /// Context lines after the matched line
    pub after: Vec<Vec<u8>>,
}

impl Match {
    fn new(lineno: usize, line: &[u8], spans: Vec<(usize, usize)>) -> Match {
        Match {
            lineno: lineno,
            line: line.into(),
            spans: spans,
            before: Vec::new(),
            after: Vec::new(),
        }
    }
}

/// Represents all matches from a single file.
#[derive(Debug)]
pub struct FileResult {
    /// File name, relative to initial argument
    pub fname: String,
    /// Is the file binary?  If yes, matches contains 0 or 1 element
    pub is_binary: bool,
    /// Do we provide (and print) context lines?
    pub has_context: bool,
    /// Matches relevant for printing
    pub matches: Vec<Match>,
}

impl FileResult {
    fn new(fname: String) -> FileResult {
        FileResult {
            fname: fname,
            is_binary: false,
            has_context: false,
            matches: Vec::new(),
        }
    }
}

/// Create a regular expression to search for matches from the given options.
///
/// The final regex is determined by several options, such as casing options
/// and options to take the search string literally.
pub fn create_rx(opts: &Opts) -> Regex {
    let mut pattern = opts.pattern.to_owned();
    if opts.literal {
        // escape regex meta-chars and create a normal pattern
        const ESCAPE: &'static str = ".?*+|^$(){}[]\\";
        pattern = pattern.chars().map(|c| {
            if ESCAPE.find(c).is_some() {
                format!("\\{}", c)
            } else {
                format!("{}", c)
            }
        }).collect();
    }
    if let Casing::Insensitive = opts.casing {
        pattern = format!("(?i){}", pattern);
    } else if let Casing::Smart = opts.casing {
        // smart casing: only case-insensitive when pattern contains no uppercase
        if !pattern.chars().any(|c| c.is_uppercase()) {
            pattern = format!("(?i){}", pattern);
        }
    }
    Regex::new(&pattern).unwrap()
}

/// Return normalized path: get rid of leading ./ and make leading // into /.
fn normalized_path(path: &Path) -> String {
    let s = path.to_string_lossy();
    if s.starts_with("./") {
        String::from(&s[2..])
    } else if s.starts_with("//") {
        String::from(&s[1..])
    } else {
        s.into_owned()
    }
}

/// Check file for binary-ness.
///
/// Currently only null-bytes are recognized to constitute binary file content.
/// However, this clashes with UTF-16 and UTF-32, so a more clever heuristic is
/// required at some point.
fn is_binary(buf: &[u8], len: usize) -> bool {
    if len == 0 {
        return false;
    }
    if len >= 3 && &buf[0..3] == b"\xEF\xBB\xBF" {
        // UTF-8 BOM
        return false;
    }
    let n = min(512, len);
    for b in buf[..n].iter() {
        if *b == b'\x00' {
            return true;  // null byte always means binary
        }
    }
    false
}

/// Iterator-like struct for collecting lines within a u8 buffer.
struct Lines<'a> {
    buf: &'a [u8],
    lines: Vec<&'a [u8]>,
    start: usize,
    lineno: usize,
}

impl<'a> Lines<'a> {
    pub fn new(buf: &[u8]) -> Lines {
        Lines { buf: buf, lines: Vec::new(), start: 0, lineno: 0 }
    }

    /// Get next line in the main iteration.
    pub fn next(&mut self) -> Option<(usize, &'a [u8])> {
        let lno = self.lineno;
        self.lineno += 1;
        if lno < self.lines.len() {
            Some((lno, self.lines[lno]))
        } else if self.advance(lno) {
            Some((lno, self.lines[lno]))
        } else {
            None
        }
    }

    /// Get an arbitrary line as a string.
    pub fn get_line(&mut self, lineno: usize) -> Option<Vec<u8>> {
        if self.advance(lineno) {
            Some(self.lines[lineno].to_vec())
        } else {
            None
        }
    }

    /// Advance the line detection until we have at least need_idx+1 lines.
    /// Return false if EOF was reached before given number of lines.
    fn advance(&mut self, need_idx: usize) -> bool {
        while self.lines.len() < need_idx + 1 {
            if self.start >= self.buf.len() {
                return false;
            }
            let end = self.buf[self.start..].iter()
                                            .position(|&x| x == b'\n')
                                            .unwrap_or(self.buf.len() - self.start);
            let line = &self.buf[self.start..self.start+end];
            self.lines.push(line);
            self.start += end + 1;
        }
        true
    }
}

/// Search a single file (represented as a u8 buffer) for matching lines.
pub fn search(regex: &Regex, opts: &Opts, path: &Path, buf: &[u8]) -> FileResult {
    let len = buf.len();
    let mut result = FileResult::new(normalized_path(path));
    result.has_context = opts.before > 0 || opts.after > 0;
    // binary file?
    if is_binary(buf, len) {
        result.is_binary = true;
        // if we care for binaries at all
        if opts.do_binaries {
            if regex.is_match(buf) {
                // found a match: create a dummy match object, and
                // leave it there (we never need more info than
                // "matched" or "didn't match")
                result.matches.push(Match::new(0, &[], Vec::new()));
            }
        }
    } else {
        let mut lines = Lines::new(buf);
        while let Some((lineno, line)) = lines.next() {
            let mut spans = Vec::new();
            if let Some(span) = regex.find(line) {
                let mut searchfrom = span.1;
                // create a match object for this line (lineno is 1-based)
                spans.push(span);
                // search for further matches in this line
                while let Some((i0, i1)) = regex.find(&line[searchfrom..]) {
                    spans.push((searchfrom + i0, searchfrom + i1));
                    searchfrom += i1;
                }
            }
            if opts.invert != spans.is_empty() {
                // no match
                continue;
            }
            let mut m = Match::new(lineno + 1, line, spans);

            // collect "before" context for this match
            if opts.before > 0 {
                for lno in lineno.saturating_sub(opts.before)..lineno {
                    m.before.push(lines.get_line(lno).unwrap());
                }
            }
            // collect "after" context for this match
            if opts.after > 0 {
                for lno in lineno+1..lineno+opts.after+1 {
                    if let Some(line) = lines.get_line(lno) {
                        m.after.push(line);
                    }
                }
            }
            result.matches.push(m);

            if opts.only_files.is_some() {
                // need only one match per file for this mode
                break;
            } else if result.matches.len() >= opts.max_count {
                break;
            }
        }
    }
    result
}
<commit_msg>Simplify Lines a bit.<commit_after>// ---------------------------------------------------------------------------------------
// Ruthenium, an ack-like searcher, (c) 2015 Georg Brandl.
// Licensed under the MIT license.
// ---------------------------------------------------------------------------------------

use std::cmp::min;
use std::path::Path;

#[cfg(feature = "pcre")]
use pcre::Regex;
#[cfg(not(feature = "pcre"))]
use regex::bytes::Regex;

use options::{Casing, Opts};

/// Represents a line that matched the pattern (maybe multiple times).
#[derive(Debug)]
pub struct Match {
    /// Line number in the file
    pub lineno: usize,
    /// Line text
    pub line: Vec<u8>,
    /// Spans (start, end) of matching parts in the line
    pub spans: Vec<(usize, usize)>,
    /// Context lines before the matched line
    pub before: Vec<Vec<u8>>,
    /// Context lines after the matched line
    pub after: Vec<Vec<u8>>,
}

impl Match {
    fn new(lineno: usize, line: &[u8], spans: Vec<(usize, usize)>) -> Match {
        Match {
            lineno: lineno,
            line: line.into(),
            spans: spans,
            before: Vec::new(),
            after: Vec::new(),
        }
    }
}

/// Represents all matches from a single file.
#[derive(Debug)]
pub struct FileResult {
    /// File name, relative to initial argument
    pub fname: String,
    /// Is the file binary?  If yes, matches contains 0 or 1 element
    pub is_binary: bool,
    /// Do we provide (and print) context lines?
    pub has_context: bool,
    /// Matches relevant for printing
    pub matches: Vec<Match>,
}

impl FileResult {
    fn new(fname: String) -> FileResult {
        FileResult {
            fname: fname,
            is_binary: false,
            has_context: false,
            matches: Vec::new(),
        }
    }
}

/// Create a regular expression to search for matches from the given options.
///
/// The final regex is determined by several options, such as casing options
/// and options to take the search string literally.
pub fn create_rx(opts: &Opts) -> Regex {
    let mut pattern = opts.pattern.to_owned();
    if opts.literal {
        // escape regex meta-chars and create a normal pattern
        const ESCAPE: &'static str = ".?*+|^$(){}[]\\";
        pattern = pattern.chars().map(|c| {
            if ESCAPE.find(c).is_some() {
                format!("\\{}", c)
            } else {
                format!("{}", c)
            }
        }).collect();
    }
    if let Casing::Insensitive = opts.casing {
        pattern = format!("(?i){}", pattern);
    } else if let Casing::Smart = opts.casing {
        // smart casing: only case-insensitive when pattern contains no uppercase
        if !pattern.chars().any(|c| c.is_uppercase()) {
            pattern = format!("(?i){}", pattern);
        }
    }
    Regex::new(&pattern).unwrap()
}

/// Return normalized path: get rid of leading ./ and make leading // into /.
fn normalized_path(path: &Path) -> String {
    let s = path.to_string_lossy();
    if s.starts_with("./") {
        String::from(&s[2..])
    } else if s.starts_with("//") {
        String::from(&s[1..])
    } else {
        s.into_owned()
    }
}

/// Check file for binary-ness.
///
/// Currently only null-bytes are recognized to constitute binary file content.
/// However, this clashes with UTF-16 and UTF-32, so a more clever heuristic is
/// required at some point.
fn is_binary(buf: &[u8], len: usize) -> bool {
    if len == 0 {
        return false;
    }
    if len >= 3 && &buf[0..3] == b"\xEF\xBB\xBF" {
        // UTF-8 BOM
        return false;
    }
    let n = min(512, len);
    for b in buf[..n].iter() {
        if *b == b'\x00' {
            return true;  // null byte always means binary
        }
    }
    false
}

/// Iterator-like struct for collecting lines within a u8 buffer.
struct Lines<'a> {
    buf: &'a [u8],
    lines: Vec<&'a [u8]>,
    lineno: usize,
}

impl<'a> Lines<'a> {
    pub fn new(buf: &[u8]) -> Lines {
        Lines { buf: buf, lines: Vec::new(), lineno: 0 }
    }

    /// Get next line in the main iteration.
    pub fn next(&mut self) -> Option<(usize, &'a [u8])> {
        let lno = self.lineno;
        self.lineno += 1;
        if lno < self.lines.len() {
            Some((lno, self.lines[lno]))
        } else if self.advance(lno) {
            Some((lno, self.lines[lno]))
        } else {
            None
        }
    }

    /// Get an arbitrary line as a string.
    pub fn get_line(&mut self, lineno: usize) -> Option<Vec<u8>> {
        if self.advance(lineno) {
            Some(self.lines[lineno].to_vec())
        } else {
            None
        }
    }

    /// Advance the line detection until we have at least need_idx+1 lines.
    /// Return false if EOF was reached before given number of lines.
    fn advance(&mut self, need_idx: usize) -> bool {
        while self.lines.len() < need_idx + 1 {
            if self.buf.len() == 0 {
                return false;
            }
            let mut split = self.buf.splitn(2, |&x| x == b'\n');
            let line = split.next().unwrap();
            self.buf = split.next().unwrap_or(&[]);
            self.lines.push(line);
        }
        true
    }
}

/// Search a single file (represented as a u8 buffer) for matching lines.
pub fn search(regex: &Regex, opts: &Opts, path: &Path, buf: &[u8]) -> FileResult {
    let len = buf.len();
    let mut result = FileResult::new(normalized_path(path));
    result.has_context = opts.before > 0 || opts.after > 0;
    // binary file?
    if is_binary(buf, len) {
        result.is_binary = true;
        // if we care for binaries at all
        if opts.do_binaries {
            if regex.is_match(buf) {
                // found a match: create a dummy match object, and
                // leave it there (we never need more info than
                // "matched" or "didn't match")
                result.matches.push(Match::new(0, &[], Vec::new()));
            }
        }
    } else {
        let mut lines = Lines::new(buf);
        while let Some((lineno, line)) = lines.next() {
            let mut spans = Vec::new();
            if let Some(span) = regex.find(line) {
                let mut searchfrom = span.1;
                // create a match object for this line (lineno is 1-based)
                spans.push(span);
                // search for further matches in this line
                while let Some((i0, i1)) = regex.find(&line[searchfrom..]) {
                    spans.push((searchfrom + i0, searchfrom + i1));
                    searchfrom += i1;
                }
            }
            if opts.invert != spans.is_empty() {
                // no match
                continue;
            }
            let mut m = Match::new(lineno + 1, line, spans);

            // collect "before" context for this match
            if opts.before > 0 {
                for lno in lineno.saturating_sub(opts.before)..lineno {
                    m.before.push(lines.get_line(lno).unwrap());
                }
            }
            // collect "after" context for this match
            if opts.after > 0 {
                for lno in lineno+1..lineno+opts.after+1 {
                    if let Some(line) = lines.get_line(lno) {
                        m.after.push(line);
                    }
                }
            }
            result.matches.push(m);

            if opts.only_files.is_some() {
                // need only one match per file for this mode
                break;
            } else if result.matches.len() >= opts.max_count {
                break;
            }
        }
    }
    result
}
<|endoftext|>


<commit_before>use nom::{multispace, space};
use nom::{IResult, Err, ErrorKind, Needed};
use std::str;

use common::{field_list, unsigned_number, statement_terminator, table_list};
use parser::ConditionExpression;

use condition::*;

#[derive(Clone, Debug, PartialEq)]
pub struct GroupByClause {
    columns: Vec<String>,
    having: String, // XXX(malte): should this be an arbitrary expr?
}

#[derive(Clone, Debug, PartialEq)]
pub struct LimitClause {
    limit: u64,
    offset: u64,
}

#[derive(Clone, Debug, PartialEq)]
pub enum OrderType {
    OrderAscending,
    OrderDescending,
}

#[derive(Clone, Debug, PartialEq)]
pub struct OrderClause {
    cols: Vec<String>, // TODO(malte): can this be an arbitrary expr?
    order: OrderType,
}

#[derive(Clone, Debug, Default, PartialEq)]
pub struct SelectStatement {
    pub tables: Vec<String>,
    pub distinct: bool,
    pub fields: Vec<String>,
    pub where_clause: Option<ConditionExpression>,
    pub group_by: Option<GroupByClause>,
    pub order: Option<OrderClause>,
    pub limit: Option<LimitClause>,
}

/// Parse LIMIT clause
named!(limit_clause<&[u8], LimitClause>,
    chain!(
        space ~
        caseless_tag!("limit") ~
        space ~
        limit_val: unsigned_number ~
        space? ~
        offset_val: opt!(chain!(
                caseless_tag!("offset") ~
                space ~
                val: unsigned_number,
                || { val }
            )
        ),
    || {
        LimitClause {
            limit: limit_val,
            offset: match offset_val {
                None => 0,
                Some(v) => v,
            },
        }
    })
);

named!(order_clause<&[u8], OrderClause>,
    dbg_dmp!(chain!(
        multispace? ~
        caseless_tag!("order by") ~
        multispace ~
        order_expr: field_list ~
        ordering: opt!(
            alt_complete!(
                  map!(caseless_tag!("desc"), |_| OrderType::OrderDescending)
                | map!(caseless_tag!("asc"), |_| OrderType::OrderAscending)
            )
        ),
    || {
        OrderClause {
            cols: order_expr.iter().map(|e| String::from(*e)).collect(),
            order: match ordering {
                None => OrderType::OrderAscending,
                Some(ref o) => o.clone(),
            },
        }
    })
));

/// Parse WHERE clause of a selection
named!(where_clause<&[u8], ConditionExpression>,
    dbg_dmp!(chain!(
        multispace? ~
        caseless_tag!("where") ~
        multispace ~
        cond: condition_expr,
        || { cond }
    )
));

/// Parse rule for a SQL selection query.
/// TODO(malte): support nested queries as selection targets
named!(pub selection<&[u8], SelectStatement>,
    chain!(
        caseless_tag!("select") ~
        space ~
        distinct: opt!(caseless_tag!("distinct")) ~
        space? ~
        fields: field_list ~
        delimited!(opt!(multispace), caseless_tag!("from"), opt!(multispace)) ~
        tables: table_list ~
        cond: opt!(where_clause) ~
        order: opt!(order_clause) ~
        limit: opt!(limit_clause) ~
        statement_terminator,
        || {
            SelectStatement {
                tables: tables.iter().map(|t| String::from(*t)).collect(),
                distinct: distinct.is_some(),
                fields: fields.iter().map(|s| String::from(*s)).collect(),
                where_clause: cond,
                group_by: None,
                order: order,
                limit: limit,
            }
        }
    )
);

#[cfg(test)]
mod tests {
    use super::*;
    use common::Operator;
    use parser::{ConditionBase, ConditionExpression, ConditionTree};

    #[test]
    fn simple_select() {
        let qstring = "SELECT id, name FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![String::from("users")],
                       fields: vec!["id".into(), "name".into()],
                       ..Default::default()
                   });
    }

    #[test]
    fn select_all() {
        let qstring = "SELECT * FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![String::from("users")],
                       fields: vec!["ALL".into()],
                       ..Default::default()
                   });
    }

    #[test]
    fn spaces_optional() {
        let qstring = "SELECT id,name FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![String::from("users")],
                       fields: vec!["id".into(), "name".into()],
                       ..Default::default()
                   });
    }

    #[test]
    fn case_sensitivity() {
        let qstring_lc = "select id, name from users;";
        let qstring_uc = "SELECT id, name FROM users;";

        assert_eq!(selection(qstring_lc.as_bytes()).unwrap(),
                   selection(qstring_uc.as_bytes()).unwrap());
    }

    #[test]
    fn termination() {
        let qstring_sem = "select id, name from users;";
        let qstring_linebreak = "select id, name from users\n";

        assert_eq!(selection(qstring_sem.as_bytes()).unwrap(),
                   selection(qstring_linebreak.as_bytes()).unwrap());
    }

    #[test]
    fn where_clause() {
        let qstring = "select * from ContactInfo where email=?;";

        let res = selection(qstring.as_bytes());

        let expected_where_cond = Some(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(String::from("email"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        }));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![String::from("ContactInfo")],
                       fields: vec!["ALL".into()],
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }

    #[test]
    fn limit_clause() {
        let qstring1 = "select * from users limit 10\n";
        let qstring2 = "select * from users limit 10 offset 10\n";

        let expected_lim1 = LimitClause {
            limit: 10,
            offset: 0,
        };
        let expected_lim2 = LimitClause {
            limit: 10,
            offset: 10,
        };

        let res1 = selection(qstring1.as_bytes());
        let res2 = selection(qstring2.as_bytes());
        assert_eq!(res1.unwrap().1.limit, Some(expected_lim1));
        assert_eq!(res2.unwrap().1.limit, Some(expected_lim2));
    }

    #[test]
    fn order_clause() {
        let qstring1 = "select * from users order by name desc\n";
        let qstring2 = "select * from users order by name, age desc\n";
        let qstring3 = "select * from users order by name\n";

        let expected_ord1 = OrderClause {
            cols: vec!["name".into()],
            order: OrderType::OrderDescending,
        };
        let expected_ord2 = OrderClause {
            cols: vec!["name".into(), "age".into()],
            order: OrderType::OrderDescending,
        };
        let expected_ord3 = OrderClause {
            cols: vec!["name".into()],
            order: OrderType::OrderAscending,
        };

        let res1 = selection(qstring1.as_bytes());
        let res2 = selection(qstring2.as_bytes());
        let res3 = selection(qstring3.as_bytes());
        assert_eq!(res1.unwrap().1.order, Some(expected_ord1));
        assert_eq!(res2.unwrap().1.order, Some(expected_ord2));
        assert_eq!(res3.unwrap().1.order, Some(expected_ord3));
    }

    #[test]
    fn table_alias() {
        let qstring1 = "select * from PaperTag as t;";
        // let qstring2 = "select * from PaperTag t;";

        let res1 = selection(qstring1.as_bytes());
        assert_eq!(res1.clone().unwrap().1,
                   SelectStatement {
                       tables: vec![String::from("PaperTag")],
                       fields: vec!["ALL".into()],
                       ..Default::default()
                   });
        // let res2 = selection(qstring2.as_bytes());
        // assert_eq!(res1.unwrap().1, res2.unwrap().1);
    }

    #[test]
    fn distinct() {
        let qstring = "select distinct tag from PaperTag where paperId=?;";

        let res = selection(qstring.as_bytes());
        let expected_where_cond = Some(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(String::from("paperId"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        }));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![String::from("PaperTag")],
                       distinct: true,
                       fields: vec!["tag".into()],
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }

    #[test]
    fn simple_condition_expr() {
        let qstring = "select infoJson from PaperStorage where paperId=? and paperStorageId=?;";

        let res = selection(qstring.as_bytes());

        let left_comp = Some(Box::new(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(String::from("paperId"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: String::from("="),
        })));
        let right_comp = Some(Box::new(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(String::from("paperStorageId"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: String::from("="),
        })));
        let expected_where_cond = Some(ConditionExpression::LogicalOp(ConditionTree {
            left: left_comp,
            right: right_comp,
            operator: String::from("and"),
        }));
        println!("res: {:#?}", res);
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![String::from("PaperStorage")],
                       fields: vec!["infoJson".into()],
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }
}
<commit_msg>Fix broken operator syntax<commit_after>use nom::{multispace, space};
use nom::{IResult, Err, ErrorKind, Needed};
use std::str;

use common::{field_list, unsigned_number, statement_terminator, table_list};
use parser::ConditionExpression;

use condition::*;

#[derive(Clone, Debug, PartialEq)]
pub struct GroupByClause {
    columns: Vec<String>,
    having: String, // XXX(malte): should this be an arbitrary expr?
}

#[derive(Clone, Debug, PartialEq)]
pub struct LimitClause {
    limit: u64,
    offset: u64,
}

#[derive(Clone, Debug, PartialEq)]
pub enum OrderType {
    OrderAscending,
    OrderDescending,
}

#[derive(Clone, Debug, PartialEq)]
pub struct OrderClause {
    cols: Vec<String>, // TODO(malte): can this be an arbitrary expr?
    order: OrderType,
}

#[derive(Clone, Debug, Default, PartialEq)]
pub struct SelectStatement {
    pub tables: Vec<String>,
    pub distinct: bool,
    pub fields: Vec<String>,
    pub where_clause: Option<ConditionExpression>,
    pub group_by: Option<GroupByClause>,
    pub order: Option<OrderClause>,
    pub limit: Option<LimitClause>,
}

/// Parse LIMIT clause
named!(limit_clause<&[u8], LimitClause>,
    chain!(
        space ~
        caseless_tag!("limit") ~
        space ~
        limit_val: unsigned_number ~
        space? ~
        offset_val: opt!(chain!(
                caseless_tag!("offset") ~
                space ~
                val: unsigned_number,
                || { val }
            )
        ),
    || {
        LimitClause {
            limit: limit_val,
            offset: match offset_val {
                None => 0,
                Some(v) => v,
            },
        }
    })
);

named!(order_clause<&[u8], OrderClause>,
    dbg_dmp!(chain!(
        multispace? ~
        caseless_tag!("order by") ~
        multispace ~
        order_expr: field_list ~
        ordering: opt!(
            alt_complete!(
                  map!(caseless_tag!("desc"), |_| OrderType::OrderDescending)
                | map!(caseless_tag!("asc"), |_| OrderType::OrderAscending)
            )
        ),
    || {
        OrderClause {
            cols: order_expr.iter().map(|e| String::from(*e)).collect(),
            order: match ordering {
                None => OrderType::OrderAscending,
                Some(ref o) => o.clone(),
            },
        }
    })
));

/// Parse WHERE clause of a selection
named!(where_clause<&[u8], ConditionExpression>,
    dbg_dmp!(chain!(
        multispace? ~
        caseless_tag!("where") ~
        multispace ~
        cond: condition_expr,
        || { cond }
    )
));

/// Parse rule for a SQL selection query.
/// TODO(malte): support nested queries as selection targets
named!(pub selection<&[u8], SelectStatement>,
    chain!(
        caseless_tag!("select") ~
        space ~
        distinct: opt!(caseless_tag!("distinct")) ~
        space? ~
        fields: field_list ~
        delimited!(opt!(multispace), caseless_tag!("from"), opt!(multispace)) ~
        tables: table_list ~
        cond: opt!(where_clause) ~
        order: opt!(order_clause) ~
        limit: opt!(limit_clause) ~
        statement_terminator,
        || {
            SelectStatement {
                tables: tables.iter().map(|t| String::from(*t)).collect(),
                distinct: distinct.is_some(),
                fields: fields.iter().map(|s| String::from(*s)).collect(),
                where_clause: cond,
                group_by: None,
                order: order,
                limit: limit,
            }
        }
    )
);

#[cfg(test)]
mod tests {
    use super::*;
    use common::Operator;
    use parser::{ConditionBase, ConditionExpression, ConditionTree};

    #[test]
    fn simple_select() {
        let qstring = "SELECT id, name FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![String::from("users")],
                       fields: vec!["id".into(), "name".into()],
                       ..Default::default()
                   });
    }

    #[test]
    fn select_all() {
        let qstring = "SELECT * FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![String::from("users")],
                       fields: vec!["ALL".into()],
                       ..Default::default()
                   });
    }

    #[test]
    fn spaces_optional() {
        let qstring = "SELECT id,name FROM users;";

        let res = selection(qstring.as_bytes());
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![String::from("users")],
                       fields: vec!["id".into(), "name".into()],
                       ..Default::default()
                   });
    }

    #[test]
    fn case_sensitivity() {
        let qstring_lc = "select id, name from users;";
        let qstring_uc = "SELECT id, name FROM users;";

        assert_eq!(selection(qstring_lc.as_bytes()).unwrap(),
                   selection(qstring_uc.as_bytes()).unwrap());
    }

    #[test]
    fn termination() {
        let qstring_sem = "select id, name from users;";
        let qstring_linebreak = "select id, name from users\n";

        assert_eq!(selection(qstring_sem.as_bytes()).unwrap(),
                   selection(qstring_linebreak.as_bytes()).unwrap());
    }

    #[test]
    fn where_clause() {
        let qstring = "select * from ContactInfo where email=?;";

        let res = selection(qstring.as_bytes());

        let expected_where_cond = Some(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(String::from("email"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        }));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![String::from("ContactInfo")],
                       fields: vec!["ALL".into()],
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }

    #[test]
    fn limit_clause() {
        let qstring1 = "select * from users limit 10\n";
        let qstring2 = "select * from users limit 10 offset 10\n";

        let expected_lim1 = LimitClause {
            limit: 10,
            offset: 0,
        };
        let expected_lim2 = LimitClause {
            limit: 10,
            offset: 10,
        };

        let res1 = selection(qstring1.as_bytes());
        let res2 = selection(qstring2.as_bytes());
        assert_eq!(res1.unwrap().1.limit, Some(expected_lim1));
        assert_eq!(res2.unwrap().1.limit, Some(expected_lim2));
    }

    #[test]
    fn order_clause() {
        let qstring1 = "select * from users order by name desc\n";
        let qstring2 = "select * from users order by name, age desc\n";
        let qstring3 = "select * from users order by name\n";

        let expected_ord1 = OrderClause {
            cols: vec!["name".into()],
            order: OrderType::OrderDescending,
        };
        let expected_ord2 = OrderClause {
            cols: vec!["name".into(), "age".into()],
            order: OrderType::OrderDescending,
        };
        let expected_ord3 = OrderClause {
            cols: vec!["name".into()],
            order: OrderType::OrderAscending,
        };

        let res1 = selection(qstring1.as_bytes());
        let res2 = selection(qstring2.as_bytes());
        let res3 = selection(qstring3.as_bytes());
        assert_eq!(res1.unwrap().1.order, Some(expected_ord1));
        assert_eq!(res2.unwrap().1.order, Some(expected_ord2));
        assert_eq!(res3.unwrap().1.order, Some(expected_ord3));
    }

    #[test]
    fn table_alias() {
        let qstring1 = "select * from PaperTag as t;";
        // let qstring2 = "select * from PaperTag t;";

        let res1 = selection(qstring1.as_bytes());
        assert_eq!(res1.clone().unwrap().1,
                   SelectStatement {
                       tables: vec![String::from("PaperTag")],
                       fields: vec!["ALL".into()],
                       ..Default::default()
                   });
        // let res2 = selection(qstring2.as_bytes());
        // assert_eq!(res1.unwrap().1, res2.unwrap().1);
    }

    #[test]
    fn distinct() {
        let qstring = "select distinct tag from PaperTag where paperId=?;";

        let res = selection(qstring.as_bytes());
        let expected_where_cond = Some(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(String::from("paperId"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        }));
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![String::from("PaperTag")],
                       distinct: true,
                       fields: vec!["tag".into()],
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }

    #[test]
    fn simple_condition_expr() {
        let qstring = "select infoJson from PaperStorage where paperId=? and paperStorageId=?;";

        let res = selection(qstring.as_bytes());

        let left_comp = Some(Box::new(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(String::from("paperId"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        })));
        let right_comp = Some(Box::new(ConditionExpression::ComparisonOp(ConditionTree {
            left:
                Some(Box::new(ConditionExpression::Base(ConditionBase::Field(String::from("paperStorageId"))))),
            right: Some(Box::new(ConditionExpression::Base(ConditionBase::Placeholder))),
            operator: Operator::Equal,
        })));
        let expected_where_cond = Some(ConditionExpression::LogicalOp(ConditionTree {
            left: left_comp,
            right: right_comp,
            operator: Operator::And,
        }));
        println!("res: {:#?}", res);
        assert_eq!(res.unwrap().1,
                   SelectStatement {
                       tables: vec![String::from("PaperStorage")],
                       fields: vec!["infoJson".into()],
                       where_clause: expected_where_cond,
                       ..Default::default()
                   });
    }
}
<|endoftext|>


<commit_before>use backend::{AnimeEntry, AnimeInfo, Status, SyncBackend};
use chrono::{Date, Local};
use error::SeriesError;
use input::{self, Answer};
use process;
use regex::Regex;
use std;
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use toml;

#[derive(Debug)]
pub struct Series<B>
where
    B: SyncBackend,
{
    sync_backend: B,
    pub episode_data: EpisodeData,
    pub save_data: SaveData,
    pub save_path: PathBuf,
}

impl<B> Series<B>
where
    B: SyncBackend,
{
    pub const DATA_FILE_NAME: &'static str = ".anup";

    pub fn from_data(episode_data: EpisodeData, sync_backend: B) -> Result<Series<B>, SeriesError> {
        let save_path = PathBuf::from(&episode_data.path).join(Series::<B>::DATA_FILE_NAME);
        let save_data = SaveData::from_path_or_default(&save_path)?;

        let series = Series {
            sync_backend,
            episode_data,
            save_data,
            save_path,
        };

        Ok(series)
    }

    pub fn load_season(&mut self, season: u32) -> Result<Season<B>, SeriesError> {
        let season_info = self.get_season_info(season)?;
        let season_ep_offset = self.calculate_season_offset(season);

        let list_entry = self.get_list_entry(season_info.clone())?;

        Season::init(self, list_entry, season_ep_offset)
    }

    pub fn save_data(&self) -> Result<(), SeriesError> {
        self.save_data.write_to(&self.save_path)
    }

    fn get_season_info(&mut self, season: u32) -> Result<AnimeInfo, SeriesError> {
        let num_seasons = self.save_data.seasons.len() as u32;

        if season >= num_seasons {
            let mut series = None;

            for cur_season in num_seasons..=season {
                println!(
                    "select the correct series for season {} of [{}]:",
                    1 + cur_season,
                    self.episode_data.series_name
                );

                let series_info = self.search_and_select_series(&self.episode_data.series_name)?;
                self.save_data.seasons.push(series_info.clone().into());

                series = Some(series_info);
            }

            self.save_data()?;

            // This unwrap should never fail, as the enclosing if statement ensures the for loop will set the
            // series variable at least once
            Ok(series.unwrap())
        } else {
            let season_info = &self.save_data.seasons[season as usize];

            let series = self
                .sync_backend
                .get_series_info_by_id(season_info.series_id)?;

            Ok(series)
        }
    }

    fn calculate_season_offset(&self, season: u32) -> u32 {
        let mut offset = 0;

        for cur_season in 0..(season as usize) {
            if let Some(season_eps) = self.save_data.seasons[cur_season].episodes {
                offset += season_eps;
            }
        }

        offset
    }

    fn search_and_select_series(&self, name: &str) -> Result<AnimeInfo, SeriesError> {
        let mut found = self.sync_backend.search_by_name(name)?;

        println!("{} results for [{}]:", B::name(), name);
        println!("enter the number next to the desired series:\n");

        println!("0 [custom search]");

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

        let index = input::read_range(0, found.len())?;

        if index == 0 {
            println!("enter the name you want to search for:");

            let name = input::read_line()?;
            self.search_and_select_series(&name)
        } else {
            let info = found.swap_remove(index - 1);
            Ok(info)
        }
    }

    fn get_list_entry(&self, info: AnimeInfo) -> Result<AnimeEntry, SeriesError> {
        let found = self.sync_backend.get_list_entry(info.clone())?;

        match found {
            Some(entry) => Ok(entry),
            None => {
                let mut entry = AnimeEntry::new(info);
                entry.status = Status::Watching;
                entry.start_date = Some(Local::today());

                self.sync_backend.update_list_entry(&entry)?;
                Ok(entry)
            }
        }
    }
}

#[derive(Debug)]
pub struct Season<'a, B>
where
    B: 'a + SyncBackend,
{
    series: &'a Series<B>,
    pub list_entry: AnimeEntry,
    pub ep_offset: u32,
}

impl<'a, B> Season<'a, B>
where
    B: 'a + SyncBackend,
{
    fn init(
        series: &'a Series<B>,
        list_entry: AnimeEntry,
        ep_offset: u32,
    ) -> Result<Season<'a, B>, SeriesError> {
        let mut season = Season {
            series,
            list_entry,
            ep_offset,
        };

        match season.list_entry.status {
            Status::Completed => season.prompt_to_rewatch()?,
            _ => (),
        }

        Ok(season)
    }

    fn update_list_entry(&self) -> Result<(), SeriesError> {
        self.series
            .sync_backend
            .update_list_entry(&self.list_entry)?;

        Ok(())
    }

    pub fn play_episode(&mut self, relative_ep: u32) -> Result<(), SeriesError> {
        let ep_num = self.ep_offset + relative_ep;

        let path = self
            .series
            .episode_data
            .episodes
            .get(&ep_num)
            .ok_or_else(|| SeriesError::EpisodeNotFound(ep_num))?;

        let status = process::open_with_default(path).map_err(SeriesError::FailedToOpenPlayer)?;
        self.list_entry.watched_episodes = relative_ep;

        if !status.success() {
            println!("video player not exited normally");
            println!("do you still want to count the episode as watched? (y/N)");

            if !input::read_yn(Answer::No)? {
                return Ok(());
            }
        }

        match self.list_entry.info.episodes {
            Some(total_eps) if relative_ep >= total_eps => self.series_completed()?,
            _ => self.episode_completed()?,
        }

        Ok(())
    }

    pub fn play_all_episodes(&mut self) -> Result<(), SeriesError> {
        loop {
            let next_ep = self.list_entry.watched_episodes + 1;

            self.play_episode(next_ep)?;
            self.next_episode_options()?;
        }
    }

    fn episode_completed(&mut self) -> Result<(), SeriesError> {
        println!(
            "[{}] episode {}/{} completed",
            self.list_entry.info.title,
            self.list_entry.watched_episodes,
            self.list_entry
                .info
                .episodes
                .map(|e| e.to_string())
                .unwrap_or_else(|| "?".to_string())
        );

        if self.list_entry.status != Status::Rewatching {
            self.list_entry.status = Status::Watching;

            if self.list_entry.watched_episodes <= 1 {
                self.list_entry.start_date = Some(Local::today());
            }
        }

        self.update_list_entry()?;
        Ok(())
    }

    fn series_completed(&mut self) -> Result<(), SeriesError> {
        println!(
            "[{}] completed!\ndo you want to rate it? (Y/n)",
            self.list_entry.info.title
        );

        if input::read_yn(Answer::Yes)? {
            self.prompt_to_update_score();
        }

        self.list_entry.status = Status::Completed;
        self.add_series_finish_date(Local::today())?;
        self.update_list_entry()?;

        // Nothing to do now
        std::process::exit(0);
    }

    fn next_episode_options(&mut self) -> Result<(), SeriesError> {
        println!("options:");
        println!("\t[d] drop series\n\t[h] put series on hold\n\t[r] rate series\n\t[x] exit\n\t[n] watch next episode (default)");

        let input = input::read_line()?.to_lowercase();

        match input.as_str() {
            "d" => {
                self.list_entry.status = Status::Dropped;
                self.add_series_finish_date(Local::today())?;
                self.update_list_entry()?;

                std::process::exit(0);
            }
            "h" => {
                self.list_entry.status = Status::OnHold;
                self.update_list_entry()?;

                std::process::exit(0);
            }
            "r" => {
                self.prompt_to_update_score();
                self.update_list_entry()?;

                self.next_episode_options()?;
            }
            "x" => std::process::exit(0),
            _ => (),
        }

        Ok(())
    }

    fn prompt_to_update_score(&mut self) {
        let max_score = self.series.sync_backend.max_score();
        println!("enter your score between 1-{}", max_score);

        match input::read_range(1.0, f32::from(max_score)) {
            Ok(score) => self.list_entry.score = score,
            Err(err) => eprintln!("failed to get score: {}", err),
        }
    }

    fn add_series_finish_date(&mut self, date: Date<Local>) -> Result<(), SeriesError> {
        let entry = &mut self.list_entry;

        // Someone may want to keep the original start / finish date for an
        // anime they're rewatching
        if entry.status == Status::Rewatching && entry.finish_date.is_some() {
            println!("do you want to override the finish date? (Y/n)");

            if input::read_yn(Answer::Yes)? {
                entry.finish_date = Some(date);
            }
        } else {
            entry.finish_date = Some(date);
        }

        Ok(())
    }

    fn prompt_to_rewatch(&mut self) -> Result<(), SeriesError> {
        println!("[{}] already completed", self.list_entry.info.title);
        println!("do you want to rewatch it? (Y/n)");

        if input::read_yn(Answer::Yes)? {
            self.list_entry.status = Status::Rewatching;
            self.list_entry.watched_episodes = 0;

            println!("do you want to reset the start and end date? (Y/n)");

            if input::read_yn(Answer::Yes)? {
                self.list_entry.start_date = Some(Local::today());
                self.list_entry.finish_date = None;
            }

            self.update_list_entry()?;
        } else {
            // No point in continuing in this case
            std::process::exit(0);
        }

        Ok(())
    }
}

type SeriesName = String;
type EpisodeNum = u32;

#[derive(Debug)]
pub struct EpisodeData {
    pub series_name: String,
    pub episodes: HashMap<u32, PathBuf>,
    pub path: PathBuf,
}

impl EpisodeData {
    pub const EP_FORMAT_REGEX: &'static str =
        r"(?:\[.+?\]\s*)?(?P<series>.+?)\s*-\s*(?P<episode>\d+).*?\..+?";

    pub fn parse_dir(dir: &Path) -> Result<EpisodeData, SeriesError> {
        if !dir.is_dir() {
            return Err(SeriesError::NotADirectory(
                dir.file_name()
                    .map(|n| n.to_string_lossy().into_owned())
                    .unwrap_or_else(|| "err".into()),
            ));
        }

        let mut series_name = None;
        let mut episodes = HashMap::new();

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

            if !path.is_file() {
                continue;
            }

            match EpisodeData::parse_filename(&path) {
                Ok((ep_name, ep_num)) => {
                    match series_name {
                        Some(ref series_name) if &ep_name != series_name => {
                            return Err(SeriesError::MultipleSeriesFound);
                        }
                        Some(_) => (),
                        None => series_name = Some(ep_name),
                    }

                    episodes.insert(ep_num, path);
                }
                Err(SeriesError::EpisodeRegexCaptureFailed) => continue,
                Err(e) => return Err(e),
            }
        }

        let name = series_name.ok_or(SeriesError::NoEpisodesFound)?;

        Ok(EpisodeData {
            series_name: name,
            episodes,
            path: dir.into(),
        })
    }

    fn parse_filename(path: &Path) -> Result<(SeriesName, EpisodeNum), SeriesError> {
        lazy_static! {
            static ref EP_FORMAT: Regex = Regex::new(EpisodeData::EP_FORMAT_REGEX).unwrap();
        }

        // Replace certain characters with spaces since they can prevent proper series
        // identification or prevent it from being found on a sync backend
        let filename = path
            .file_name()
            .and_then(|path| path.to_str())
            .ok_or(SeriesError::UnableToGetFilename)?
            .replace('_', " ");

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

        let series_name = {
            let raw_name = &caps["series"];

            raw_name.replace('.', " ")
            .replace(" -", ":") // Dashes typically represent a colon in file names
            .trim()
            .to_string()
        };

        let episode = caps["episode"]
            .parse()
            .map_err(SeriesError::EpisodeNumParseFailed)?;

        Ok((series_name, episode))
    }
}

#[derive(Debug, Default, Serialize, Deserialize)]
pub struct SaveData {
    pub seasons: Vec<SeasonInfo>,
}

impl SaveData {
    fn from_path(path: &Path) -> Result<SaveData, SeriesError> {
        let file_contents = fs::read_to_string(path)?;
        let data = toml::from_str(&file_contents)?;

        Ok(data)
    }

    fn from_path_or_default(path: &Path) -> Result<SaveData, SeriesError> {
        if path.exists() {
            SaveData::from_path(path)
        } else {
            Ok(SaveData::default())
        }
    }

    fn write_to(&self, path: &Path) -> Result<(), SeriesError> {
        let toml = toml::to_string_pretty(self)?;
        fs::write(path, toml)?;

        Ok(())
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct SeasonInfo {
    pub series_id: u32,
    pub episodes: Option<u32>,
}

impl From<AnimeInfo> for SeasonInfo {
    fn from(info: AnimeInfo) -> SeasonInfo {
        SeasonInfo {
            series_id: info.id,
            episodes: info.episodes,
        }
    }
}
<commit_msg>Refactor the way list entry properties get updated<commit_after>use backend::{AnimeEntry, AnimeInfo, Status, SyncBackend};
use chrono::Local;
use error::SeriesError;
use input::{self, Answer};
use process;
use regex::Regex;
use std;
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use toml;

#[derive(Debug)]
pub struct Series<B>
where
    B: SyncBackend,
{
    sync_backend: B,
    pub episode_data: EpisodeData,
    pub save_data: SaveData,
    pub save_path: PathBuf,
}

impl<B> Series<B>
where
    B: SyncBackend,
{
    pub const DATA_FILE_NAME: &'static str = ".anup";

    pub fn from_data(episode_data: EpisodeData, sync_backend: B) -> Result<Series<B>, SeriesError> {
        let save_path = PathBuf::from(&episode_data.path).join(Series::<B>::DATA_FILE_NAME);
        let save_data = SaveData::from_path_or_default(&save_path)?;

        let series = Series {
            sync_backend,
            episode_data,
            save_data,
            save_path,
        };

        Ok(series)
    }

    pub fn load_season(&mut self, season: u32) -> Result<Season<B>, SeriesError> {
        let season_info = self.get_season_info(season)?;
        let season_ep_offset = self.calculate_season_offset(season);

        let list_entry = self.get_list_entry(season_info.clone())?;

        Season::init(self, list_entry, season_ep_offset)
    }

    pub fn save_data(&self) -> Result<(), SeriesError> {
        self.save_data.write_to(&self.save_path)
    }

    fn get_season_info(&mut self, season: u32) -> Result<AnimeInfo, SeriesError> {
        let num_seasons = self.save_data.seasons.len() as u32;

        if season >= num_seasons {
            let mut series = None;

            for cur_season in num_seasons..=season {
                println!(
                    "select the correct series for season {} of [{}]:",
                    1 + cur_season,
                    self.episode_data.series_name
                );

                let series_info = self.search_and_select_series(&self.episode_data.series_name)?;
                self.save_data.seasons.push(series_info.clone().into());

                series = Some(series_info);
            }

            self.save_data()?;

            // This unwrap should never fail, as the enclosing if statement ensures the for loop will set the
            // series variable at least once
            Ok(series.unwrap())
        } else {
            let season_info = &self.save_data.seasons[season as usize];

            let series = self
                .sync_backend
                .get_series_info_by_id(season_info.series_id)?;

            Ok(series)
        }
    }

    fn calculate_season_offset(&self, season: u32) -> u32 {
        let mut offset = 0;

        for cur_season in 0..(season as usize) {
            if let Some(season_eps) = self.save_data.seasons[cur_season].episodes {
                offset += season_eps;
            }
        }

        offset
    }

    fn search_and_select_series(&self, name: &str) -> Result<AnimeInfo, SeriesError> {
        let mut found = self.sync_backend.search_by_name(name)?;

        println!("{} results for [{}]:", B::name(), name);
        println!("enter the number next to the desired series:\n");

        println!("0 [custom search]");

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

        let index = input::read_range(0, found.len())?;

        if index == 0 {
            println!("enter the name you want to search for:");

            let name = input::read_line()?;
            self.search_and_select_series(&name)
        } else {
            let info = found.swap_remove(index - 1);
            Ok(info)
        }
    }

    fn get_list_entry(&self, info: AnimeInfo) -> Result<AnimeEntry, SeriesError> {
        let found = self.sync_backend.get_list_entry(info.clone())?;

        match found {
            Some(entry) => Ok(entry),
            None => {
                let mut entry = AnimeEntry::new(info);
                entry.status = Status::Watching;
                entry.start_date = Some(Local::today());

                self.sync_backend.update_list_entry(&entry)?;
                Ok(entry)
            }
        }
    }
}

#[derive(Debug)]
pub struct Season<'a, B>
where
    B: 'a + SyncBackend,
{
    series: &'a Series<B>,
    pub list_entry: AnimeEntry,
    pub ep_offset: u32,
}

impl<'a, B> Season<'a, B>
where
    B: 'a + SyncBackend,
{
    fn init(
        series: &'a Series<B>,
        list_entry: AnimeEntry,
        ep_offset: u32,
    ) -> Result<Season<'a, B>, SeriesError> {
        let mut season = Season {
            series,
            list_entry,
            ep_offset,
        };

        match season.list_entry.status {
            Status::Completed => season.update_series_status(Status::Rewatching)?,
            _ => (),
        }

        Ok(season)
    }

    fn update_list_entry(&self) -> Result<(), SeriesError> {
        self.series
            .sync_backend
            .update_list_entry(&self.list_entry)?;

        Ok(())
    }

    pub fn play_episode(&mut self, relative_ep: u32) -> Result<(), SeriesError> {
        let ep_num = self.ep_offset + relative_ep;

        let path = self
            .series
            .episode_data
            .episodes
            .get(&ep_num)
            .ok_or_else(|| SeriesError::EpisodeNotFound(ep_num))?;

        let status = process::open_with_default(path).map_err(SeriesError::FailedToOpenPlayer)?;
        self.list_entry.watched_episodes = relative_ep;

        if !status.success() {
            println!("video player not exited normally");
            println!("do you still want to count the episode as watched? (y/N)");

            if !input::read_yn(Answer::No)? {
                return Ok(());
            }
        }

        match self.list_entry.info.episodes {
            Some(total_eps) if relative_ep >= total_eps => {
                self.update_series_status(Status::Completed)?;
                // TODO: replace with more graceful way of exiting
                std::process::exit(0);
            }
            _ => self.episode_completed()?,
        }

        Ok(())
    }

    pub fn play_all_episodes(&mut self) -> Result<(), SeriesError> {
        loop {
            let next_ep = self.list_entry.watched_episodes + 1;

            self.play_episode(next_ep)?;
            self.next_episode_options()?;
        }
    }

    fn episode_completed(&self) -> Result<(), SeriesError> {
        println!(
            "[{}] episode {}/{} completed",
            self.list_entry.info.title,
            self.list_entry.watched_episodes,
            self.list_entry
                .info
                .episodes
                .map(|e| e.to_string())
                .unwrap_or_else(|| "?".to_string())
        );

        self.update_list_entry()
    }

    fn next_episode_options(&mut self) -> Result<(), SeriesError> {
        println!("options:");
        println!("\t[d] drop series\n\t[h] put series on hold\n\t[r] rate series\n\t[x] exit\n\t[n] watch next episode (default)");

        let input = input::read_line()?.to_lowercase();

        match input.as_str() {
            "d" => {
                self.update_series_status(Status::Dropped)?;
                // TODO: replace with more graceful way of exiting
                std::process::exit(0);
            }
            "h" => {
                self.update_series_status(Status::OnHold)?;
                // TODO: replace with more graceful way of exiting
                std::process::exit(0);
            }
            "r" => {
                self.prompt_to_update_score();
                self.update_list_entry()?;

                self.next_episode_options()?;
            }
            "x" => std::process::exit(0),
            _ => (),
        }

        Ok(())
    }

    fn prompt_to_update_score(&mut self) {
        let max_score = self.series.sync_backend.max_score();
        println!("enter your score between 1-{}", max_score);

        match input::read_range(1.0, f32::from(max_score)) {
            Ok(score) => self.list_entry.score = score,
            Err(err) => eprintln!("failed to get score: {}", err),
        }
    }

    fn update_series_status(&mut self, status: Status) -> Result<(), SeriesError> {
        self.list_entry.status = status;

        match status {
            Status::Watching => {
                if self.list_entry.watched_episodes <= 1 {
                    self.list_entry.start_date = Some(Local::today());
                }
            }
            Status::Rewatching => {
                println!("[{}] already completed", self.list_entry.info.title);
                println!("do you want to reset the start and end dates of the series? (Y/n)");

                if input::read_yn(Answer::Yes)? {
                    self.list_entry.start_date = Some(Local::today());
                    self.list_entry.finish_date = None;
                }

                self.list_entry.watched_episodes = 0;
            }
            Status::Completed => {
                if self.list_entry.finish_date.is_none() {
                    self.list_entry.finish_date = Some(Local::today());
                }

                println!(
                    "[{}] completed!\ndo you want to rate it? (Y/n)",
                    self.list_entry.info.title
                );

                if input::read_yn(Answer::Yes)? {
                    self.prompt_to_update_score();
                }
            }
            Status::Dropped => {
                if self.list_entry.finish_date.is_none() {
                    self.list_entry.finish_date = Some(Local::today());
                }
            }
            Status::OnHold | Status::PlanToWatch => (),
        }

        self.update_list_entry()?;
        Ok(())
    }
}

type SeriesName = String;
type EpisodeNum = u32;

#[derive(Debug)]
pub struct EpisodeData {
    pub series_name: String,
    pub episodes: HashMap<u32, PathBuf>,
    pub path: PathBuf,
}

impl EpisodeData {
    pub const EP_FORMAT_REGEX: &'static str =
        r"(?:\[.+?\]\s*)?(?P<series>.+?)\s*-\s*(?P<episode>\d+).*?\..+?";

    pub fn parse_dir(dir: &Path) -> Result<EpisodeData, SeriesError> {
        if !dir.is_dir() {
            return Err(SeriesError::NotADirectory(
                dir.file_name()
                    .map(|n| n.to_string_lossy().into_owned())
                    .unwrap_or_else(|| "err".into()),
            ));
        }

        let mut series_name = None;
        let mut episodes = HashMap::new();

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

            if !path.is_file() {
                continue;
            }

            match EpisodeData::parse_filename(&path) {
                Ok((ep_name, ep_num)) => {
                    match series_name {
                        Some(ref series_name) if &ep_name != series_name => {
                            return Err(SeriesError::MultipleSeriesFound);
                        }
                        Some(_) => (),
                        None => series_name = Some(ep_name),
                    }

                    episodes.insert(ep_num, path);
                }
                Err(SeriesError::EpisodeRegexCaptureFailed) => continue,
                Err(e) => return Err(e),
            }
        }

        let name = series_name.ok_or(SeriesError::NoEpisodesFound)?;

        Ok(EpisodeData {
            series_name: name,
            episodes,
            path: dir.into(),
        })
    }

    fn parse_filename(path: &Path) -> Result<(SeriesName, EpisodeNum), SeriesError> {
        lazy_static! {
            static ref EP_FORMAT: Regex = Regex::new(EpisodeData::EP_FORMAT_REGEX).unwrap();
        }

        // Replace certain characters with spaces since they can prevent proper series
        // identification or prevent it from being found on a sync backend
        let filename = path
            .file_name()
            .and_then(|path| path.to_str())
            .ok_or(SeriesError::UnableToGetFilename)?
            .replace('_', " ");

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

        let series_name = {
            let raw_name = &caps["series"];

            raw_name.replace('.', " ")
            .replace(" -", ":") // Dashes typically represent a colon in file names
            .trim()
            .to_string()
        };

        let episode = caps["episode"]
            .parse()
            .map_err(SeriesError::EpisodeNumParseFailed)?;

        Ok((series_name, episode))
    }
}

#[derive(Debug, Default, Serialize, Deserialize)]
pub struct SaveData {
    pub seasons: Vec<SeasonInfo>,
}

impl SaveData {
    fn from_path(path: &Path) -> Result<SaveData, SeriesError> {
        let file_contents = fs::read_to_string(path)?;
        let data = toml::from_str(&file_contents)?;

        Ok(data)
    }

    fn from_path_or_default(path: &Path) -> Result<SaveData, SeriesError> {
        if path.exists() {
            SaveData::from_path(path)
        } else {
            Ok(SaveData::default())
        }
    }

    fn write_to(&self, path: &Path) -> Result<(), SeriesError> {
        let toml = toml::to_string_pretty(self)?;
        fs::write(path, toml)?;

        Ok(())
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct SeasonInfo {
    pub series_id: u32,
    pub episodes: Option<u32>,
}

impl From<AnimeInfo> for SeasonInfo {
    fn from(info: AnimeInfo) -> SeasonInfo {
        SeasonInfo {
            series_id: info.id,
            episodes: info.episodes,
        }
    }
}
<|endoftext|>


<commit_before>use message;

pub struct Server {
    id: i32,
}

static CONNACK_OK : [u8; 4] = [32, 2, 0, 0];
static PING_RESP : [u8; 2] = [0xd0, 0];

impl Server {
    pub fn new() -> Self {
        Server { id: 4 }
    }

    pub fn new_message<C: Client>(&mut self, client: &mut C, bytes: &[u8]) {
        let message_type = message::message_type(bytes);
        match message_type {
            message::MqttType::Connect => {
                client.send(&CONNACK_OK);
            },
            message::MqttType::PingReq => {
                client.send(&PING_RESP);
            },
            _ => panic!("Unknown message type")
        }
    }

    pub fn new_bytes<C: Client>(&mut self, client: &mut C, bytes: &[u8]) {
        let mut slice = bytes;

        while slice.len() > 0 {
            let remaining_len = message::remaining_length(bytes);
            let total_len = remaining_len + 2;

            if slice.len() >= total_len {
                let msg = &slice[0 .. total_len];
                slice = &slice[total_len ..];
                self.new_message(client, msg);
            } else {
                slice = &[];
            }
        }
    }
}

struct Stream {
    buffer: Vec<u8>,
    bytes_start: usize,
}

impl Stream {
    pub fn new() -> Self {
        Stream { buffer: vec![0; 1024], bytes_start: 0 }
    }

    pub fn buffer(&mut self) -> &mut Vec<u8> {
        &mut self.buffer
    }

    pub fn handle_messages<C: Client>(&mut self, bytes_read: usize, server: &mut Server, client: &mut C) {
        let mut slice = &self.buffer[self.bytes_start .. self.bytes_start + bytes_read];
        let mut total_len = 2;
        while slice.len() >= total_len {
            total_len = message::remaining_length(slice) + 2; //2 for fixed header size
            let msg = &slice[0 .. total_len];
            slice = &slice[total_len..];
            server.new_message(client, msg);
        }
    }
}

pub trait Client {
    fn send(&mut self, bytes: &[u8]);
}

struct TestClient {
    msgs: Vec<Vec<u8>>,
}

impl TestClient {
    fn new() -> Self {
        TestClient { msgs: vec!() }
    }

    fn last_msg(&self) -> &[u8] {
        self.msgs.last().unwrap()
    }

    fn read(&self, buffer: &mut [u8], bytes: &[u8]) -> usize {
        for i in 0..bytes.len() {
            buffer[i] = bytes[i];
        }
        bytes.len()
    }
}

impl Client for TestClient {
    fn send(&mut self, bytes: &[u8]) {
        self.msgs.push(bytes.to_vec());
    }
}

#[test]
fn test_connect() {
    let connect_bytes = &[
        0x10u8, 0x2a, // fixed header
        0x00, 0x06, 'M' as u8, 'Q' as u8, 'I' as u8, 's' as u8, 'd' as u8, 'p' as u8,
        0x03, // protocol version
        0xcc, // connection flags 1100111x user, pw, !wr, w(01), w, !c, x
        0x00, 0x0a, // keepalive of 100
        0x00, 0x03, 'c' as u8, 'i' as u8, 'd' as u8, // client ID
        0x00, 0x04, 'w' as u8, 'i' as u8, 'l' as u8, 'l' as u8, // will topic
        0x00, 0x04, 'w' as u8, 'm' as u8, 's' as u8, 'g' as u8, // will msg
        0x00, 0x07, 'g' as u8, 'l' as u8, 'i' as u8, 'f' as u8, 't' as u8, 'e' as u8, 'l' as u8, // username
        0x00, 0x02, 'p' as u8, 'w' as u8, // password
        ][0..];

    let mut server = Server::new();
    let mut client = TestClient::new();

    server.new_message(&mut client, connect_bytes);
    assert_eq!(client.last_msg(), &CONNACK_OK);
}

#[test]
fn test_ping() {
    let ping_bytes =  &[0xc0u8, 0][0..];

    let mut server = Server::new();
    let mut client = TestClient::new();

    server.new_message(&mut client, ping_bytes);

    assert_eq!(client.last_msg(), &PING_RESP);
}

#[test]
fn test_pings_all_at_once() {
    let ping_bytes = &[0xc0u8, 0, 0xc0, 0, 0xc0, 0, 0xc0, 0][0..];

    let mut server = Server::new();
    let mut client = TestClient::new();
    let mut stream = Stream::new();

    let bytes_read = client.read(stream.buffer(), ping_bytes);
    stream.handle_messages(bytes_read, &mut server, &mut client);

    assert_eq!(client.msgs.len(), 4);
    for msg in client.msgs {
        assert_eq!(msg, &PING_RESP);
    }
}

#[test]
fn test_pings_multiple_time_unbroken() {
    let ping_bytes = &[0xc0u8, 0, 0xc0, 0][0..];

    let mut server = Server::new();
    let mut client = TestClient::new();
    let mut stream = Stream::new();

    let bytes_read = client.read(stream.buffer(), ping_bytes);
    stream.handle_messages(bytes_read, &mut server, &mut client);
    assert_eq!(client.msgs.len(), 2);

    let bytes_read = client.read(stream.buffer(), ping_bytes);
    stream.handle_messages(bytes_read, &mut server, &mut client);
    assert_eq!(client.msgs.len(), 4);

    for msg in client.msgs {
        assert_eq!(msg, &PING_RESP);
    }
}

#[test]
fn test_pings_broken() {
    let ping_fst = &[0xc0u8][0..];
    let ping_snd = &[0u8][0..];

    let mut server = Server::new();
    let mut client = TestClient::new();
    let mut stream = Stream::new();

    let bytes_read = client.read(stream.buffer(), ping_fst);
    stream.handle_messages(bytes_read, &mut server, &mut client);
    assert_eq!(client.msgs.len(), 0);

    let bytes_read = client.read(stream.buffer(), ping_snd);
    stream.handle_messages(bytes_read, &mut server, &mut client);
    assert_eq!(client.msgs.len(), 1);

    let bytes_read = client.read(stream.buffer(), ping_fst);
    stream.handle_messages(bytes_read, &mut server, &mut client);
    assert_eq!(client.msgs.len(), 1);

    let bytes_read = client.read(stream.buffer(), ping_snd);
    stream.handle_messages(bytes_read, &mut server, &mut client);
    assert_eq!(client.msgs.len(), 2);

    for msg in client.msgs {
        assert_eq!(msg, &PING_RESP);
    }
}
<commit_msg>UT for broken message passes, but can't shift bytes back to the beginning of the buffer<commit_after>use message;

pub struct Server {
    id: i32,
}

static CONNACK_OK : [u8; 4] = [32, 2, 0, 0];
static PING_RESP : [u8; 2] = [0xd0, 0];

impl Server {
    pub fn new() -> Self {
        Server { id: 4 }
    }

    pub fn new_message<C: Client>(&mut self, client: &mut C, bytes: &[u8]) {
        let message_type = message::message_type(bytes);
        match message_type {
            message::MqttType::Connect => {
                client.send(&CONNACK_OK);
            },
            message::MqttType::PingReq => {
                client.send(&PING_RESP);
            },
            _ => panic!("Unknown message type")
        }
    }

    pub fn new_bytes<C: Client>(&mut self, client: &mut C, bytes: &[u8]) {
        let mut slice = bytes;

        while slice.len() > 0 {
            let remaining_len = message::remaining_length(bytes);
            let total_len = remaining_len + 2;

            if slice.len() >= total_len {
                let msg = &slice[0 .. total_len];
                slice = &slice[total_len ..];
                self.new_message(client, msg);
            } else {
                slice = &[];
            }
        }
    }
}

struct Stream {
    //the reason there's bytes_start and bytes_read is because bytes_read always grows,
    //but bytes_start only moves if full messages have been processed
    buffer: Vec<u8>,
    bytes_start: usize, //the start of the next byte window
    bytes_read: usize,  //bytes read so far
}

impl Stream {
    pub fn new() -> Self {
        Stream { buffer: vec![0; 1024], bytes_start: 0, bytes_read: 0 }
    }

    pub fn buffer(&mut self) -> &mut [u8] {
        &mut self.buffer[self.bytes_read .. ]
    }

    pub fn handle_messages<C: Client>(&mut self, bytes_read: usize, server: &mut Server, client: &mut C) {
        let mut slice = &self.buffer[self.bytes_start .. self.bytes_read + bytes_read];
        const HEADER_LEN: usize = 2;
        let mut total_len = HEADER_LEN;
        while slice.len() >= total_len {
            total_len = message::remaining_length(slice) + HEADER_LEN;
            let msg = &slice[0 .. total_len];
            slice = &slice[total_len..];
            self.bytes_start += total_len;
            server.new_message(client, msg);
        }

        // //shift everything to the beginning of the buffer
        // for i in 0..slice.len() {
        //     self.buffer[i] = slice[i];
        // }

        self.bytes_read += bytes_read;
    }
}

pub trait Client {
    fn send(&mut self, bytes: &[u8]);
}

struct TestClient {
    msgs: Vec<Vec<u8>>,
}

impl TestClient {
    fn new() -> Self {
        TestClient { msgs: vec!() }
    }

    fn last_msg(&self) -> &[u8] {
        self.msgs.last().unwrap()
    }

    fn read(&self, buffer: &mut [u8], bytes: &[u8]) -> usize {
        for i in 0..bytes.len() {
            buffer[i] = bytes[i];
        }
        bytes.len()
    }
}

impl Client for TestClient {
    fn send(&mut self, bytes: &[u8]) {
        self.msgs.push(bytes.to_vec());
    }
}

#[test]
fn test_connect() {
    let connect_bytes = &[
        0x10u8, 0x2a, // fixed header
        0x00, 0x06, 'M' as u8, 'Q' as u8, 'I' as u8, 's' as u8, 'd' as u8, 'p' as u8,
        0x03, // protocol version
        0xcc, // connection flags 1100111x user, pw, !wr, w(01), w, !c, x
        0x00, 0x0a, // keepalive of 100
        0x00, 0x03, 'c' as u8, 'i' as u8, 'd' as u8, // client ID
        0x00, 0x04, 'w' as u8, 'i' as u8, 'l' as u8, 'l' as u8, // will topic
        0x00, 0x04, 'w' as u8, 'm' as u8, 's' as u8, 'g' as u8, // will msg
        0x00, 0x07, 'g' as u8, 'l' as u8, 'i' as u8, 'f' as u8, 't' as u8, 'e' as u8, 'l' as u8, // username
        0x00, 0x02, 'p' as u8, 'w' as u8, // password
        ][0..];

    let mut server = Server::new();
    let mut client = TestClient::new();

    server.new_message(&mut client, connect_bytes);
    assert_eq!(client.last_msg(), &CONNACK_OK);
}

#[test]
fn test_ping() {
    let ping_bytes =  &[0xc0u8, 0][0..];

    let mut server = Server::new();
    let mut client = TestClient::new();

    server.new_message(&mut client, ping_bytes);

    assert_eq!(client.last_msg(), &PING_RESP);
}

#[test]
fn test_pings_all_at_once() {
    let ping_bytes = &[0xc0u8, 0, 0xc0, 0, 0xc0, 0, 0xc0, 0][0..];

    let mut server = Server::new();
    let mut client = TestClient::new();
    let mut stream = Stream::new();

    let bytes_read = client.read(stream.buffer(), ping_bytes);
    stream.handle_messages(bytes_read, &mut server, &mut client);

    assert_eq!(client.msgs.len(), 4);
    for msg in client.msgs {
        assert_eq!(msg, &PING_RESP);
    }
}

#[test]
fn test_pings_multiple_time_unbroken() {
    let ping_bytes = &[0xc0u8, 0, 0xc0, 0][0..];

    let mut server = Server::new();
    let mut client = TestClient::new();
    let mut stream = Stream::new();

    let bytes_read = client.read(stream.buffer(), ping_bytes);
    stream.handle_messages(bytes_read, &mut server, &mut client);
    assert_eq!(client.msgs.len(), 2);

    let bytes_read = client.read(stream.buffer(), ping_bytes);
    stream.handle_messages(bytes_read, &mut server, &mut client);
    assert_eq!(client.msgs.len(), 4);

    for msg in client.msgs {
        assert_eq!(msg, &PING_RESP);
    }
}

#[test]
fn test_pings_broken() {
    let ping_fst = &[0xc0u8][0..];
    let ping_snd = &[0u8][0..];

    let mut server = Server::new();
    let mut client = TestClient::new();
    let mut stream = Stream::new();

    let bytes_read = client.read(stream.buffer(), ping_fst);
    stream.handle_messages(bytes_read, &mut server, &mut client);
    assert_eq!(client.msgs.len(), 0);

    let bytes_read = client.read(stream.buffer(), ping_snd);
    stream.handle_messages(bytes_read, &mut server, &mut client);
    assert_eq!(client.msgs.len(), 1);

    let bytes_read = client.read(stream.buffer(), ping_fst);
    stream.handle_messages(bytes_read, &mut server, &mut client);
    assert_eq!(client.msgs.len(), 1);

    let bytes_read = client.read(stream.buffer(), ping_snd);
    stream.handle_messages(bytes_read, &mut server, &mut client);
    assert_eq!(client.msgs.len(), 2);

    for msg in client.msgs {
        assert_eq!(msg, &PING_RESP);
    }
}
<|endoftext|>