<commit_before>use curl::easy::{Easy, List};
use json;
use metrics::Metric;
use mpmc::Queue;
use shuteye::sleep;
use std::default::Default;
use std::fmt;
use std::io::Read;
use std::process::Command;
use std::time::Duration;
use tic::{Clocksource, Sample, Sender};
use webhook::event::*;

pub struct Config {
    events: Option<Queue<Event>>,
    clock: Option<Clocksource>,
    stats: Option<Sender<Metric>>,
    token: Option<String>,
}

impl Config {
    pub fn build(self) -> Result<Consumer, &'static str> {
        Consumer::configured(self)
    }

    pub fn clock(mut self, clock: Clocksource) -> Self {
        self.clock = Some(clock);
        self
    }

    pub fn events(mut self, queue: Queue<Event>) -> Self {
        self.events = Some(queue);
        self
    }

    pub fn stats(mut self, sender: Sender<Metric>) -> Self {
        self.stats = Some(sender);
        self
    }

    pub fn token(mut self, token: String) -> Self {
        self.token = Some(token);
        self
    }
}

impl Default for Config {
    fn default() -> Config {
        Config {
            events: None,
            clock: None,
            stats: None,
            token: None,
        }
    }
}

pub enum Status {
    Pending,
    Success,
    Error,
    Failure,
}

impl fmt::Display for Status {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Status::Pending => write!(f, "pending"),
            Status::Success => write!(f, "success"),
            Status::Error => write!(f, "error"),
            Status::Failure => write!(f, "failure"),
        }
    }
}

pub struct Consumer {
    clock: Clocksource,
    stats: Sender<Metric>,
    events: Queue<Event>,
    token: String,
}

impl Consumer {
    pub fn configure() -> Config {
        Default::default()
    }

    pub fn configured(config: Config) -> Result<Consumer, &'static str> {
        let events = config.events.clone();
        let clock = config.clock.clone();
        let stats = config.stats.clone();
        let token = config.token.clone();

        if events.is_none() {
            return Err("need events queue");
        }
        if clock.is_none() {
            return Err("need clock");
        }
        if stats.is_none() {
            return Err("need stats");
        }
        if token.is_none() {
            return Err("need token");
        }
        Ok(Consumer {
               events: events.unwrap(),
               clock: clock.unwrap(),
               stats: stats.unwrap(),
               token: token.unwrap(),
           })
    }

    fn time(&self) -> u64 {
        self.clock.counter()
    }

    pub fn run(&mut self) {
        if let Some(event) = self.events.pop() {
            let t0 = self.time();
            trace!("consume event: {:?}", event);
            // do processing
            match event {
                Event::PullRequest(event) => self.handle_pull_request(event),
                Event::Push(event) => self.handle_push(event),
                _ => {}
            }
            let t1 = self.time();
            let _ = self.stats.send(Sample::new(t0, t1, Metric::Processed));
        }
        // call sleep
        sleep(Duration::new(0, 1_000_000));
    }

    fn send_status(&self,
                   repo: &str,
                   sha: &str,
                   state: &str,
                   context: &str,
                   description: &str,
                   url: &str) {
        info!("set: {} to: {}", sha, state);
        let endpoint = format!("https://api.github.com/repos/{}/statuses/{}", repo, sha);
        let auth = format!("Authorization: token {}", self.token);
        let mut list = List::new();
        list.append(&auth).unwrap();
        list.append("content-type: application/json").unwrap();

        let data = object!{
            "state" => state,
            "target_url" => url,
            "description" => description,
            "context" => context
        };

        trace!("sending: {}", data);

        let mut data_to_upload = data.dump();
        let mut handle = Easy::new();
        handle.useragent("crucible");
        handle.url(&endpoint).unwrap();
        handle.http_headers(list).unwrap();
        handle.post(true).unwrap();
        let mut response = Vec::new();

        handle.post_fields_copy(&data_to_upload.as_bytes());
        {
            let mut transfer = handle.transfer();
            transfer
                .write_function(|new_data| {
                                    response.extend_from_slice(new_data);
                                    Ok(new_data.len())
                                })
                .unwrap();
            transfer.perform().unwrap();
        }
        let rsp_string = String::from_utf8(response).unwrap_or("invalid utf8".to_owned());
        trace!("response: {}", rsp_string);
    }

    fn handle_push(&mut self, event: Push) {
        // this gets scary
        let base_path = "/mnt/scratch/";

        let id = "temp";
        let path = base_path.to_owned() + id;

        // inform github we're running a test
        self.send_status(&event.repo(),
                         &event.sha(),
                         "pending",
                         "continuous-integration/crucible/push",
                         "pending...",
                         "https://oxidize.io");

        // prepare
        create_directory(&path);
        let status = clone_repo(&path, &event.repo(), &event.url(), &event.sha());
        if status.is_err() {
            self.send_status(&event.repo(),
                             &event.sha(),
                             "error",
                             "continuous-integration/crucible/push",
                             "whoops. error.",
                             "https://oxidize.io");
        } else {
            // run test
            let result_test = cargo_test(&path);
            let result_fmt = cargo_fmt(&path);

            // this should send a real result
            if result_test.is_err() || result_fmt.is_err() {
                self.send_status(&event.repo(),
                                 &event.sha(),
                                 "failed",
                                 "continuous-integration/crucible/push",
                                 "the build failed",
                                 "https://oxidize.io");
            } else {
                self.send_status(&event.repo(),
                                 &event.sha(),
                                 "success",
                                 "continuous-integration/crucible/push",
                                 "lgtm. shipit",
                                 "https://oxidize.io");
            }
        }

        // cleanup
        remove_directory(&path);
    }

    fn handle_pull_request(&mut self, event: PullRequest) {}
}

fn clone_repo(path: &str, name: &str, url: &str, sha: &str) -> Result<(), ()> {
    info!("clone repo: {}", name);
    let status = Command::new("git")
        .arg("clone")
        .arg(url)
        .arg("repo")
        .current_dir(path)
        .status()
        .expect("failed to run git");
    if !status.success() {
        return Err(());
    }
    let status = Command::new("git")
        .arg("checkout")
        .arg(sha)
        .current_dir(path.to_owned() + "/repo")
        .status()
        .expect("failed to run git");
    if !status.success() {
        return Err(());
    }
    Ok(())
}

fn create_directory(path: &str) {
    Command::new("mkdir")
        .arg("-p")
        .arg(path)
        .status()
        .expect("failed to run mkdir");
}

fn remove_directory(path: &str) {
    Command::new("rm")
        .arg("-rf")
        .arg(path)
        .status()
        .expect("failed to run rm");
}

fn cargo_test(path: &str) -> Result<(), ()> {
    let status = Command::new("cargo")
        .arg("test")
        .current_dir(path.to_owned() + "/repo")
        .status()
        .expect("failed to run cargo test");
    if status.success() {
        info!("cargo test: passed");
        Ok(())
    } else {
        info!("cargo test: failed");
        Err(())
    }
}
//  cargo fmt -- --write-mode=diff
fn cargo_fmt(path: &str) -> Result<(), ()> {
    let status = Command::new("cargo")
        .arg("fmt")
        .arg("--")
        .arg("--write-mode=diff")
        .current_dir(path.to_owned() + "/repo")
        .status()
        .expect("failed to run cargo fmt");
    if status.success() {
        info!("cargo fmt: passed");
        Ok(())
    } else {
        info!("cargo fmt: failed");
        Err(())
    }
}
<commit_msg>bugfix: deleted branch (#15)<commit_after>use curl::easy::{Easy, List};
use json;
use metrics::Metric;
use mpmc::Queue;
use shuteye::sleep;
use std::default::Default;
use std::fmt;
use std::io::Read;
use std::process::Command;
use std::time::Duration;
use tic::{Clocksource, Sample, Sender};
use webhook::event::*;

pub struct Config {
    events: Option<Queue<Event>>,
    clock: Option<Clocksource>,
    stats: Option<Sender<Metric>>,
    token: Option<String>,
}

impl Config {
    pub fn build(self) -> Result<Consumer, &'static str> {
        Consumer::configured(self)
    }

    pub fn clock(mut self, clock: Clocksource) -> Self {
        self.clock = Some(clock);
        self
    }

    pub fn events(mut self, queue: Queue<Event>) -> Self {
        self.events = Some(queue);
        self
    }

    pub fn stats(mut self, sender: Sender<Metric>) -> Self {
        self.stats = Some(sender);
        self
    }

    pub fn token(mut self, token: String) -> Self {
        self.token = Some(token);
        self
    }
}

impl Default for Config {
    fn default() -> Config {
        Config {
            events: None,
            clock: None,
            stats: None,
            token: None,
        }
    }
}

pub enum Status {
    Pending,
    Success,
    Error,
    Failure,
}

impl fmt::Display for Status {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Status::Pending => write!(f, "pending"),
            Status::Success => write!(f, "success"),
            Status::Error => write!(f, "error"),
            Status::Failure => write!(f, "failure"),
        }
    }
}

pub struct Consumer {
    clock: Clocksource,
    stats: Sender<Metric>,
    events: Queue<Event>,
    token: String,
}

impl Consumer {
    pub fn configure() -> Config {
        Default::default()
    }

    pub fn configured(config: Config) -> Result<Consumer, &'static str> {
        let events = config.events.clone();
        let clock = config.clock.clone();
        let stats = config.stats.clone();
        let token = config.token.clone();

        if events.is_none() {
            return Err("need events queue");
        }
        if clock.is_none() {
            return Err("need clock");
        }
        if stats.is_none() {
            return Err("need stats");
        }
        if token.is_none() {
            return Err("need token");
        }
        Ok(Consumer {
               events: events.unwrap(),
               clock: clock.unwrap(),
               stats: stats.unwrap(),
               token: token.unwrap(),
           })
    }

    fn time(&self) -> u64 {
        self.clock.counter()
    }

    pub fn run(&mut self) {
        if let Some(event) = self.events.pop() {
            let t0 = self.time();
            trace!("consume event: {:?}", event);
            // do processing
            match event {
                Event::PullRequest(event) => self.handle_pull_request(event),
                Event::Push(event) => self.handle_push(event),
                _ => {}
            }
            let t1 = self.time();
            let _ = self.stats.send(Sample::new(t0, t1, Metric::Processed));
        }
        // call sleep
        sleep(Duration::new(0, 1_000_000));
    }

    fn send_status(&self,
                   repo: &str,
                   sha: &str,
                   state: &str,
                   context: &str,
                   description: &str,
                   url: &str) {
        info!("set: {} to: {}", sha, state);
        let endpoint = format!("https://api.github.com/repos/{}/statuses/{}", repo, sha);
        let auth = format!("Authorization: token {}", self.token);
        let mut list = List::new();
        list.append(&auth).unwrap();
        list.append("content-type: application/json").unwrap();

        let data = object!{
            "state" => state,
            "target_url" => url,
            "description" => description,
            "context" => context
        };

        trace!("sending: {}", data);

        let mut data_to_upload = data.dump();
        let mut handle = Easy::new();
        handle.useragent("crucible");
        handle.url(&endpoint).unwrap();
        handle.http_headers(list).unwrap();
        handle.post(true).unwrap();
        let mut response = Vec::new();

        handle.post_fields_copy(&data_to_upload.as_bytes());
        {
            let mut transfer = handle.transfer();
            transfer
                .write_function(|new_data| {
                                    response.extend_from_slice(new_data);
                                    Ok(new_data.len())
                                })
                .unwrap();
            transfer.perform().unwrap();
        }
        let rsp_string = String::from_utf8(response).unwrap_or("invalid utf8".to_owned());
        trace!("response: {}", rsp_string);
    }

    fn handle_push(&mut self, event: Push) {
        // this gets scary
        let base_path = "/mnt/scratch/";

        let id = "temp";
        let path = base_path.to_owned() + id;

        // skip events with this sha, happens when branch deleted
        if event.sha() == "0000000000000000000000000000000000000000" {
            return;
        }

        // inform github we're running a test
        self.send_status(&event.repo(),
                         &event.sha(),
                         "pending",
                         "continuous-integration/crucible/push",
                         "pending...",
                         "https://oxidize.io");

        // prepare
        create_directory(&path);
        let status = clone_repo(&path, &event.repo(), &event.url(), &event.sha());
        if status.is_err() {
            self.send_status(&event.repo(),
                             &event.sha(),
                             "error",
                             "continuous-integration/crucible/push",
                             "whoops. error.",
                             "https://oxidize.io");
        } else {
            // run test
            let result_test = cargo_test(&path);
            let result_fmt = cargo_fmt(&path);

            // this should send a real result
            if result_test.is_err() || result_fmt.is_err() {
                self.send_status(&event.repo(),
                                 &event.sha(),
                                 "failed",
                                 "continuous-integration/crucible/push",
                                 "the build failed",
                                 "https://oxidize.io");
            } else {
                self.send_status(&event.repo(),
                                 &event.sha(),
                                 "success",
                                 "continuous-integration/crucible/push",
                                 "lgtm. shipit",
                                 "https://oxidize.io");
            }
        }

        // cleanup
        remove_directory(&path);
    }

    fn handle_pull_request(&mut self, event: PullRequest) {}
}

fn clone_repo(path: &str, name: &str, url: &str, sha: &str) -> Result<(), ()> {
    info!("clone repo: {}", name);
    let status = Command::new("git")
        .arg("clone")
        .arg(url)
        .arg("repo")
        .current_dir(path)
        .status()
        .expect("failed to run git");
    if !status.success() {
        return Err(());
    }
    let status = Command::new("git")
        .arg("checkout")
        .arg(sha)
        .current_dir(path.to_owned() + "/repo")
        .status()
        .expect("failed to run git");
    if !status.success() {
        return Err(());
    }
    Ok(())
}

fn create_directory(path: &str) {
    Command::new("mkdir")
        .arg("-p")
        .arg(path)
        .status()
        .expect("failed to run mkdir");
}

fn remove_directory(path: &str) {
    Command::new("rm")
        .arg("-rf")
        .arg(path)
        .status()
        .expect("failed to run rm");
}

fn cargo_test(path: &str) -> Result<(), ()> {
    let status = Command::new("cargo")
        .arg("test")
        .current_dir(path.to_owned() + "/repo")
        .status()
        .expect("failed to run cargo test");
    if status.success() {
        info!("cargo test: passed");
        Ok(())
    } else {
        info!("cargo test: failed");
        Err(())
    }
}
//  cargo fmt -- --write-mode=diff
fn cargo_fmt(path: &str) -> Result<(), ()> {
    let status = Command::new("cargo")
        .arg("fmt")
        .arg("--")
        .arg("--write-mode=diff")
        .current_dir(path.to_owned() + "/repo")
        .status()
        .expect("failed to run cargo fmt");
    if status.success() {
        info!("cargo fmt: passed");
        Ok(())
    } else {
        info!("cargo fmt: failed");
        Err(())
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::path::Path;

use leveldb;
use leveldb::database::Database as LDatabase;
use leveldb::iterator::Iterable;
use leveldb::kv::KV;
use leveldb::options::{Options, ReadOptions, WriteOptions};
use leveldb::snapshots::Snapshots;
use db_key::Key as DBKey;

use futures::{IntoFuture, Done};

// Not quite a fan of this implementation
pub struct Key {
    data: Vec<u8>,
}

impl DBKey for Key {
    fn from_u8(key: &[u8]) -> Self {
        Key { data: key.to_vec() }
    }

    fn as_slice<T, F: Fn(&[u8]) -> T>(&self, f: F) -> T {
        f(&*self.data)
    }
}

#[derive(Debug)]
pub enum Error {
    LevelDBError(leveldb::error::Error),
    IOError(io::Error),
}

impl Error {
    fn is_level_db_error(&self) -> bool {
        match *self {
            Error::LevelDBError(_) => true,
        }
    }
}

impl PartialEq for Error {
    fn eq(&self, other: &Self) -> bool {
        self.is_level_db_error() && other.is_level_db_error()
    }
}

impl Eq for Error {}

impl From<leveldb::error::Error> for Error {
    fn from(e: leveldb::error::Error) -> Self {
        Error::LevelDBError(e)
    }
}

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self {
        Error::IOError(e)
    }
}

pub trait Database: Send + Sync + 'static {
    fn put(&self, key: &[u8], value: &[u8]) -> Done<(), Error>;
    fn get(&self, key: &[u8]) -> Done<Option<Vec<u8>>, Error>;
    fn list(&self, key: &[u8]) -> Done<Vec<Vec<u8>>, Error>;
}

pub struct LevelDB<K: DBKey> {
    db: LDatabase<K>,
}

impl<K: DBKey> LevelDB<K> {
    pub fn new(path: &str) -> Result<Self, Error> {
        let mut opts = Options::new();
        opts.create_if_missing = true;
        let db = try!(LDatabase::open(Path::new(path), opts));
        Ok(LevelDB { db: db })
    }
}

impl<K: DBKey + 'static> Database for LevelDB<K> {
    fn put(&self, key: &[u8], value: &[u8]) -> Done<(), Error> {
        let write_opts = WriteOptions::new();
        self.db.put(write_opts, K::from_u8(key), value).map_err(|e| Error::LevelDBError(e)).into_future()
    }

    fn get(&self, key: &[u8]) -> Done<Option<Vec<u8>>, Error> {
        let read_opts = ReadOptions::new();
        self.db.get(read_opts, K::from_u8(key)).map_err(|e| Error::LevelDBError(e)).into_future()
    }

    fn list(&self, key: &[u8]) -> Done<Vec<Vec<u8>>, Error> {
        let read_opts = ReadOptions::new();
        // this is awful, but rust's leveldb lib does not seem to provide filtering for us
        Ok(self.db.snapshot().iter(read_opts).filter(|k| { K::as_slice(&k.0, |x| &x[..key.len()] == key )}).map(|k| k.1).collect()).into_future()
    }
}

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

    use std::fmt;

    use futures::{Future, Task};
    use super::{Database, Key, LevelDB};
    use std::sync::Arc;

    pub fn assert_done<T, F>(mut f: F, result: Result<T::Item, T::Error>)
        where T: Future,
        T::Item: Eq + fmt::Debug,
        T::Error: Eq + fmt::Debug,
        F: FnMut() -> T,
    {
        let mut a = f();
        assert_eq!(&a.poll(&mut Task::new()).unwrap(), &result);
        drop(a);
    }

    #[test]
    fn put_get() {
        let db = Arc::new(LevelDB::<Key>::new(tempdir::TempDir::new("put_get").unwrap().path().to_str().unwrap()).unwrap());
        let key = [0, 1, 2, 3];
        let value = [4, 6, 8, 10];
        assert_done(move ||  {
            let db = db.clone();
            db.put(&key, &value).and_then(move |_| {
                db.get(&key)
            })
        }, Ok(Some(value.to_vec())))
    }

    #[test]
    fn put_put_get() {
        let db = Arc::new(LevelDB::<Key>::new(tempdir::TempDir::new("put_get").unwrap().path().to_str().unwrap()).unwrap());
        let key = [0, 1, 2, 3];
        let value = [4, 6, 8, 10];
        let value2 = [100, 101, 102, 103];
        assert_done(move ||  {
            let db = db.clone();
            db.put(&key, &value).and_then(move |_| {
                db.put(&key, &value2).and_then(move |_| {
                    db.get(&key)
                })
            })
        }, Ok(Some(value2.to_vec())))
    }

    #[test]
    fn get_none() {
        let db = LevelDB::<Key>::new(tempdir::TempDir::new("get_none").unwrap().path().to_str().unwrap()).unwrap();
        let key = [0, 1, 2, 3];
        assert_done(move ||  {
            db.get(&key)
        }, Ok(None));
    }

    #[test]
    fn put_list() {
        let db = Arc::new(LevelDB::<Key>::new(tempdir::TempDir::new("put_list").unwrap().path().to_str().unwrap()).unwrap());
        let key = [0, 1, 2, 3];
        let value = [4, 6, 8, 10];
        let key2 = [0, 1, 2, 3, 4];
        let value2 = [100, 101, 102, 103, 104];
        assert_done(move ||  {
            let db = db.clone();
            db.put(&key, &value).and_then(move |_| {
                db.put(&key2, &value2).and_then(move |_| {
                    db.list(&key)
                })
            })
        }, Ok(vec![value.to_vec(), value2.to_vec()]));
    }
}
<commit_msg>Match all cases<commit_after>use std::io;
use std::path::Path;

use leveldb;
use leveldb::database::Database as LDatabase;
use leveldb::iterator::Iterable;
use leveldb::kv::KV;
use leveldb::options::{Options, ReadOptions, WriteOptions};
use leveldb::snapshots::Snapshots;
use db_key::Key as DBKey;

use futures::{IntoFuture, Done};

// Not quite a fan of this implementation
pub struct Key {
    data: Vec<u8>,
}

impl DBKey for Key {
    fn from_u8(key: &[u8]) -> Self {
        Key { data: key.to_vec() }
    }

    fn as_slice<T, F: Fn(&[u8]) -> T>(&self, f: F) -> T {
        f(&*self.data)
    }
}

#[derive(Debug)]
pub enum Error {
    LevelDBError(leveldb::error::Error),
    IOError(io::Error),
}

impl Error {
    fn is_level_db_error(&self) -> bool {
        match *self {
            Error::LevelDBError(_) => true,
            _ => false,
        }
    }
}

impl PartialEq for Error {
    fn eq(&self, other: &Self) -> bool {
        self.is_level_db_error() && other.is_level_db_error()
    }
}

impl Eq for Error {}

impl From<leveldb::error::Error> for Error {
    fn from(e: leveldb::error::Error) -> Self {
        Error::LevelDBError(e)
    }
}

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self {
        Error::IOError(e)
    }
}

pub trait Database: Send + Sync + 'static {
    fn put(&self, key: &[u8], value: &[u8]) -> Done<(), Error>;
    fn get(&self, key: &[u8]) -> Done<Option<Vec<u8>>, Error>;
    fn list(&self, key: &[u8]) -> Done<Vec<Vec<u8>>, Error>;
}

pub struct LevelDB<K: DBKey> {
    db: LDatabase<K>,
}

impl<K: DBKey> LevelDB<K> {
    pub fn new(path: &str) -> Result<Self, Error> {
        let mut opts = Options::new();
        opts.create_if_missing = true;
        let db = try!(LDatabase::open(Path::new(path), opts));
        Ok(LevelDB { db: db })
    }
}

impl<K: DBKey + 'static> Database for LevelDB<K> {
    fn put(&self, key: &[u8], value: &[u8]) -> Done<(), Error> {
        let write_opts = WriteOptions::new();
        self.db.put(write_opts, K::from_u8(key), value).map_err(|e| Error::LevelDBError(e)).into_future()
    }

    fn get(&self, key: &[u8]) -> Done<Option<Vec<u8>>, Error> {
        let read_opts = ReadOptions::new();
        self.db.get(read_opts, K::from_u8(key)).map_err(|e| Error::LevelDBError(e)).into_future()
    }

    fn list(&self, key: &[u8]) -> Done<Vec<Vec<u8>>, Error> {
        let read_opts = ReadOptions::new();
        // this is awful, but rust's leveldb lib does not seem to provide filtering for us
        Ok(self.db.snapshot().iter(read_opts).filter(|k| { K::as_slice(&k.0, |x| &x[..key.len()] == key )}).map(|k| k.1).collect()).into_future()
    }
}

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

    use std::fmt;

    use futures::{Future, Task};
    use super::{Database, Key, LevelDB};
    use std::sync::Arc;

    pub fn assert_done<T, F>(mut f: F, result: Result<T::Item, T::Error>)
        where T: Future,
        T::Item: Eq + fmt::Debug,
        T::Error: Eq + fmt::Debug,
        F: FnMut() -> T,
    {
        let mut a = f();
        assert_eq!(&a.poll(&mut Task::new()).unwrap(), &result);
        drop(a);
    }

    #[test]
    fn put_get() {
        let db = Arc::new(LevelDB::<Key>::new(tempdir::TempDir::new("put_get").unwrap().path().to_str().unwrap()).unwrap());
        let key = [0, 1, 2, 3];
        let value = [4, 6, 8, 10];
        assert_done(move ||  {
            let db = db.clone();
            db.put(&key, &value).and_then(move |_| {
                db.get(&key)
            })
        }, Ok(Some(value.to_vec())))
    }

    #[test]
    fn put_put_get() {
        let db = Arc::new(LevelDB::<Key>::new(tempdir::TempDir::new("put_get").unwrap().path().to_str().unwrap()).unwrap());
        let key = [0, 1, 2, 3];
        let value = [4, 6, 8, 10];
        let value2 = [100, 101, 102, 103];
        assert_done(move ||  {
            let db = db.clone();
            db.put(&key, &value).and_then(move |_| {
                db.put(&key, &value2).and_then(move |_| {
                    db.get(&key)
                })
            })
        }, Ok(Some(value2.to_vec())))
    }

    #[test]
    fn get_none() {
        let db = LevelDB::<Key>::new(tempdir::TempDir::new("get_none").unwrap().path().to_str().unwrap()).unwrap();
        let key = [0, 1, 2, 3];
        assert_done(move ||  {
            db.get(&key)
        }, Ok(None));
    }

    #[test]
    fn put_list() {
        let db = Arc::new(LevelDB::<Key>::new(tempdir::TempDir::new("put_list").unwrap().path().to_str().unwrap()).unwrap());
        let key = [0, 1, 2, 3];
        let value = [4, 6, 8, 10];
        let key2 = [0, 1, 2, 3, 4];
        let value2 = [100, 101, 102, 103, 104];
        assert_done(move ||  {
            let db = db.clone();
            db.put(&key, &value).and_then(move |_| {
                db.put(&key2, &value2).and_then(move |_| {
                    db.list(&key)
                })
            })
        }, Ok(vec![value.to_vec(), value2.to_vec()]));
    }
}
<|endoftext|>


<commit_before>use core::ops::Deref;
use diesel;
use diesel::prelude::*;
use diesel::sqlite::SqliteConnection;
use regex::Regex;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time;

use db::DB;
use db::schema::{directories, songs};
use errors::*;
use metadata;
use vfs::Vfs;

const INDEX_BUILDING_INSERT_BUFFER_SIZE: usize = 1000; // Insertions in each transaction

pub struct IndexConfig {
	pub album_art_pattern: Option<Regex>,
	pub sleep_duration: u64, // in seconds
	pub path: PathBuf,
}

impl IndexConfig {
	pub fn new() -> IndexConfig {
		IndexConfig {
			sleep_duration: 60 * 30, // 30 minutes
			album_art_pattern: None,
			path: Path::new(":memory:").to_path_buf(),
		}
	}
}

#[derive(Debug, Insertable)]
#[table_name="songs"]
struct NewSong {
	path: String,
	parent: String,
	track_number: Option<i32>,
	disc_number: Option<i32>,
	title: Option<String>,
	artist: Option<String>,
	album_artist: Option<String>,
	year: Option<i32>,
	album: Option<String>,
	artwork: Option<String>,
}

#[derive(Debug, Insertable)]
#[table_name="directories"]
struct NewDirectory {
	path: String,
	parent: Option<String>,
	artist: Option<String>,
	year: Option<i32>,
	album: Option<String>,
	artwork: Option<String>,
	date_added: i32,
}

struct IndexBuilder<'db> {
	new_songs: Vec<NewSong>,
	new_directories: Vec<NewDirectory>,
	connection: &'db Mutex<SqliteConnection>,
}

impl<'db> IndexBuilder<'db> {
	fn new(connection: &Mutex<SqliteConnection>) -> Result<IndexBuilder> {
		let mut new_songs = Vec::new();
		let mut new_directories = Vec::new();
		new_songs.reserve_exact(INDEX_BUILDING_INSERT_BUFFER_SIZE);
		new_directories.reserve_exact(INDEX_BUILDING_INSERT_BUFFER_SIZE);
		Ok(IndexBuilder {
		       new_songs: new_songs,
		       new_directories: new_directories,
		       connection: connection,
		   })
	}

	fn flush_songs(&mut self) -> Result<()> {
		let connection = self.connection.lock().unwrap();
		let connection = connection.deref();
		connection
			.transaction::<_, Error, _>(|| {
				                            diesel::insert(&self.new_songs)
				                                .into(songs::table)
				                                .execute(connection)?;
				                            Ok(())
				                           })?;
		self.new_songs.clear();
		Ok(())
	}

	fn flush_directories(&mut self) -> Result<()> {
		let connection = self.connection.lock().unwrap();
		let connection = connection.deref();
		connection
			.transaction::<_, Error, _>(|| {
				                            diesel::insert(&self.new_directories)
				                                .into(directories::table)
				                                .execute(connection)?;
				                            Ok(())
				                           })?;
		self.new_directories.clear();
		Ok(())
	}

	fn push_song(&mut self, song: NewSong) -> Result<()> {
		if self.new_songs.len() >= self.new_songs.capacity() {
			self.flush_songs()?;
		}
		self.new_songs.push(song);
		Ok(())
	}

	fn push_directory(&mut self, directory: NewDirectory) -> Result<()> {
		if self.new_directories.len() >= self.new_directories.capacity() {
			self.flush_directories()?;
		}
		self.new_directories.push(directory);
		Ok(())
	}
}

pub struct Index {
	vfs: Arc<Vfs>,
	connection: Arc<Mutex<SqliteConnection>>,
	album_art_pattern: Option<Regex>,
	sleep_duration: u64,
}

impl Index {
	pub fn new(vfs: Arc<Vfs>,
	           connection: Arc<Mutex<SqliteConnection>>,
	           config: &IndexConfig)
	           -> Index {
		let index = Index {
			vfs: vfs,
			connection: connection,
			album_art_pattern: config.album_art_pattern.clone(),
			sleep_duration: config.sleep_duration,
		};
		index
	}

	pub fn update_index(&self) -> Result<()> {
		let start = time::Instant::now();
		println!("Beginning library index update");
		self.clean()?;
		self.populate()?;
		println!("Library index update took {} seconds",
		         start.elapsed().as_secs());
		Ok(())
	}

	fn clean(&self) -> Result<()> {
		{
			let all_songs: Vec<String>;
			{
				let connection = self.connection.lock().unwrap();
				let connection = connection.deref();
				all_songs = songs::table
					.select(songs::columns::path)
					.load(connection)?;
			}

			let missing_songs = all_songs
				.into_iter()
				.filter(|ref song_path| {
					        let path = Path::new(&song_path);
					        !path.exists() || self.vfs.real_to_virtual(path).is_err()
					       })
				.collect::<Vec<_>>();

			let connection = self.connection.lock().unwrap();
			let connection = connection.deref();
			diesel::delete(songs::table.filter(songs::columns::path.eq_any(missing_songs)))
				.execute(connection)?;
		}

		{
			let all_directories: Vec<String>;
			{
				let connection = self.connection.lock().unwrap();
				let connection = connection.deref();
				all_directories = directories::table
					.select(directories::columns::path)
					.load(connection)?;
			}

			let missing_directories = all_directories
				.into_iter()
				.filter(|ref directory_path| {
					        let path = Path::new(&directory_path);
					        !path.exists() || self.vfs.real_to_virtual(path).is_err()
					       })
				.collect::<Vec<_>>();

			let connection = self.connection.lock().unwrap();
			let connection = connection.deref();
			diesel::delete(directories::table.filter(directories::columns::path
			                                             .eq_any(missing_directories)))
					.execute(connection)?;
		}

		Ok(())
	}

	fn populate(&self) -> Result<()> {
		let vfs = self.vfs.deref();
		let mount_points = vfs.get_mount_points();
		let mut builder = IndexBuilder::new(&self.connection)?;
		for (_, target) in mount_points {
			self.populate_directory(&mut builder, None, target.as_path())?;
		}
		builder.flush_songs()?;
		builder.flush_directories()?;
		Ok(())
	}

	fn get_artwork(&self, dir: &Path) -> Option<String> {
		let pattern = match self.album_art_pattern {
			Some(ref p) => p,
			_ => return None,
		};

		if let Ok(dir_content) = fs::read_dir(dir) {
			for file in dir_content {
				if let Ok(file) = file {
					if let Some(name_string) = file.file_name().to_str() {
						if pattern.is_match(name_string) {
							return file.path().to_str().map(|p| p.to_owned());
						}
					}
				}
			}
		}

		None
	}

	fn populate_directory(&self,
	                      builder: &mut IndexBuilder,
	                      parent: Option<&Path>,
	                      path: &Path)
	                      -> Result<()> {

		// Find artwork
		let artwork = self.get_artwork(path);

		// Extract path and parent path
		let parent_string = parent.and_then(|p| p.to_str()).map(|s| s.to_owned());
		let path_string = path.to_str().ok_or("Invalid directory path")?;

		// Find date added
		let metadata = fs::metadata(path_string)?;
		let created = metadata
			.created()
			.or(metadata.modified())?
			.duration_since(time::UNIX_EPOCH)?
			.as_secs() as i32;

		let mut directory_album = None;
		let mut directory_year = None;
		let mut directory_artist = None;
		let mut inconsistent_directory_album = false;
		let mut inconsistent_directory_year = false;
		let mut inconsistent_directory_artist = false;

		// Insert content
		if let Ok(dir_content) = fs::read_dir(path) {
			for file in dir_content {
				let file_path = match file {
					Ok(f) => f.path(),
					_ => continue,
				};

				if file_path.is_dir() {
					self.populate_directory(builder, Some(path), file_path.as_path())?;
				} else {
					if let Some(file_path_string) = file_path.to_str() {
						if let Ok(tags) = metadata::read(file_path.as_path()) {
							if tags.year.is_some() {
								inconsistent_directory_year |= directory_year.is_some() &&
								                               directory_year != tags.year;
								directory_year = tags.year;
							}

							if tags.album.is_some() {
								inconsistent_directory_album |= directory_album.is_some() &&
								                                directory_album != tags.album;
								directory_album = tags.album.as_ref().map(|a| a.clone());
							}

							if tags.album_artist.is_some() {
								inconsistent_directory_artist |= directory_artist.is_some() &&
								                                 directory_artist !=
								                                 tags.album_artist;
								directory_artist = tags.album_artist.as_ref().map(|a| a.clone());
							} else if tags.artist.is_some() {
								inconsistent_directory_artist |= directory_artist.is_some() &&
								                                 directory_artist != tags.artist;
								directory_artist = tags.artist.as_ref().map(|a| a.clone());
							}

							let song = NewSong {
								path: file_path_string.to_owned(),
								parent: path_string.to_owned(),
								disc_number: tags.disc_number.map(|n| n as i32),
								track_number: tags.track_number.map(|n| n as i32),
								title: tags.title,
								artist: tags.artist,
								album_artist: tags.album_artist,
								album: tags.album,
								year: tags.year,
								artwork: artwork.as_ref().map(|s| s.to_owned()),
							};

							builder.push_song(song)?;
						}
					}
				}
			}
		}

		// Insert directory
		if inconsistent_directory_year {
			directory_year = None;
		}
		if inconsistent_directory_album {
			directory_album = None;
		}
		if inconsistent_directory_artist {
			directory_artist = None;
		}

		let directory = NewDirectory {
			path: path_string.to_owned(),
			parent: parent_string,
			artwork: artwork,
			album: directory_album,
			artist: directory_artist,
			year: directory_year,
			date_added: created,
		};
		builder.push_directory(directory)?;

		Ok(())
	}

	pub fn update_loop(&self) {
		loop {
			if let Err(e) = self.update_index() {
				println!("Error while updating index: {}", e);
			}
			thread::sleep(time::Duration::from_secs(self.sleep_duration));
		}
	}
}

fn _get_test_db(name: &str) -> DB {
	use vfs::VfsConfig;
	use std::collections::HashMap;

	let mut collection_path = PathBuf::new();
	collection_path.push("test");
	collection_path.push("collection");
	let mut mount_points = HashMap::new();
	mount_points.insert("root".to_owned(), collection_path);

	let vfs = Arc::new(Vfs::new(VfsConfig { mount_points: mount_points }));

	let mut index_config = IndexConfig::new();
	index_config.album_art_pattern = Some(Regex::new(r#"^Folder\.(png|jpg|jpeg)$"#).unwrap());
	index_config.path = PathBuf::new();
	index_config.path.push("test");
	index_config.path.push(name);

	if index_config.path.exists() {
		fs::remove_file(&index_config.path).unwrap();
	}

	DB::new(vfs, &index_config).unwrap()
}

#[test]
fn test_populate() {
	use db::models::*;

	let db = _get_test_db("populate.sqlite");
	let index = db.get_index();
	index.update_index().unwrap();
	index.update_index().unwrap(); // Check that subsequent updates don't run into conflicts

	let connection = index.connection.lock().unwrap();
	let connection = connection.deref();
	let all_directories: Vec<Directory> = directories::table.load(connection).unwrap();
	let all_songs: Vec<Song> = songs::table.load(connection).unwrap();
	assert_eq!(all_directories.len(), 5);
	assert_eq!(all_songs.len(), 12);
}

#[test]
fn test_metadata() {
	use db::models::*;

	let mut target = PathBuf::new();
	target.push("test");
	target.push("collection");
	target.push("Tobokegao");
	target.push("Picnic");

	let mut song_path = target.clone();
	song_path.push("05 - シャーベット (Sherbet).mp3");

	let mut artwork_path = target.clone();
	artwork_path.push("Folder.png");

	let db = _get_test_db("metadata.sqlite");
	let index = db.get_index();
	index.update_index().unwrap();

	let connection = index.connection.lock().unwrap();
	let connection = connection.deref();
	let songs: Vec<Song> = songs::table
		.filter(songs::columns::title.eq("シャーベット (Sherbet)"))
		.load(connection)
		.unwrap();

	assert_eq!(songs.len(), 1);
	let song = &songs[0];
	assert_eq!(song.path, song_path.to_string_lossy().as_ref());
	assert_eq!(song.track_number, Some(5));
	assert_eq!(song.disc_number, None);
	assert_eq!(song.title, Some("シャーベット (Sherbet)".to_owned()));
	assert_eq!(song.artist, Some("Tobokegao".to_owned()));
	assert_eq!(song.album_artist, None);
	assert_eq!(song.album, Some("Picnic".to_owned()));
	assert_eq!(song.year, Some(2016));
	assert_eq!(song.artwork,
	           Some(artwork_path.to_string_lossy().into_owned()));
}
<commit_msg>Autoformat<commit_after>use core::ops::Deref;
use diesel;
use diesel::prelude::*;
use diesel::sqlite::SqliteConnection;
use regex::Regex;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time;

use db::DB;
use db::schema::{directories, songs};
use errors::*;
use metadata;
use vfs::Vfs;

const INDEX_BUILDING_INSERT_BUFFER_SIZE: usize = 1000; // Insertions in each transaction

pub struct IndexConfig {
	pub album_art_pattern: Option<Regex>,
	pub sleep_duration: u64, // in seconds
	pub path: PathBuf,
}

impl IndexConfig {
	pub fn new() -> IndexConfig {
		IndexConfig {
			sleep_duration: 60 * 30, // 30 minutes
			album_art_pattern: None,
			path: Path::new(":memory:").to_path_buf(),
		}
	}
}

#[derive(Debug, Insertable)]
#[table_name="songs"]
struct NewSong {
	path: String,
	parent: String,
	track_number: Option<i32>,
	disc_number: Option<i32>,
	title: Option<String>,
	artist: Option<String>,
	album_artist: Option<String>,
	year: Option<i32>,
	album: Option<String>,
	artwork: Option<String>,
}

#[derive(Debug, Insertable)]
#[table_name="directories"]
struct NewDirectory {
	path: String,
	parent: Option<String>,
	artist: Option<String>,
	year: Option<i32>,
	album: Option<String>,
	artwork: Option<String>,
	date_added: i32,
}

struct IndexBuilder<'db> {
	new_songs: Vec<NewSong>,
	new_directories: Vec<NewDirectory>,
	connection: &'db Mutex<SqliteConnection>,
}

impl<'db> IndexBuilder<'db> {
	fn new(connection: &Mutex<SqliteConnection>) -> Result<IndexBuilder> {
		let mut new_songs = Vec::new();
		let mut new_directories = Vec::new();
		new_songs.reserve_exact(INDEX_BUILDING_INSERT_BUFFER_SIZE);
		new_directories.reserve_exact(INDEX_BUILDING_INSERT_BUFFER_SIZE);
		Ok(IndexBuilder {
		       new_songs: new_songs,
		       new_directories: new_directories,
		       connection: connection,
		   })
	}

	fn flush_songs(&mut self) -> Result<()> {
		let connection = self.connection.lock().unwrap();
		let connection = connection.deref();
		connection
			.transaction::<_, Error, _>(|| {
				                            diesel::insert(&self.new_songs)
				                                .into(songs::table)
				                                .execute(connection)?;
				                            Ok(())
				                           })?;
		self.new_songs.clear();
		Ok(())
	}

	fn flush_directories(&mut self) -> Result<()> {
		let connection = self.connection.lock().unwrap();
		let connection = connection.deref();
		connection
			.transaction::<_, Error, _>(|| {
				                            diesel::insert(&self.new_directories)
				                                .into(directories::table)
				                                .execute(connection)?;
				                            Ok(())
				                           })?;
		self.new_directories.clear();
		Ok(())
	}

	fn push_song(&mut self, song: NewSong) -> Result<()> {
		if self.new_songs.len() >= self.new_songs.capacity() {
			self.flush_songs()?;
		}
		self.new_songs.push(song);
		Ok(())
	}

	fn push_directory(&mut self, directory: NewDirectory) -> Result<()> {
		if self.new_directories.len() >= self.new_directories.capacity() {
			self.flush_directories()?;
		}
		self.new_directories.push(directory);
		Ok(())
	}
}

pub struct Index {
	vfs: Arc<Vfs>,
	connection: Arc<Mutex<SqliteConnection>>,
	album_art_pattern: Option<Regex>,
	sleep_duration: u64,
}

impl Index {
	pub fn new(vfs: Arc<Vfs>,
	           connection: Arc<Mutex<SqliteConnection>>,
	           config: &IndexConfig)
	           -> Index {
		let index = Index {
			vfs: vfs,
			connection: connection,
			album_art_pattern: config.album_art_pattern.clone(),
			sleep_duration: config.sleep_duration,
		};
		index
	}

	pub fn update_index(&self) -> Result<()> {
		let start = time::Instant::now();
		println!("Beginning library index update");
		self.clean()?;
		self.populate()?;
		println!("Library index update took {} seconds",
		         start.elapsed().as_secs());
		Ok(())
	}

	fn clean(&self) -> Result<()> {
		{
			let all_songs: Vec<String>;
			{
				let connection = self.connection.lock().unwrap();
				let connection = connection.deref();
				all_songs = songs::table
					.select(songs::columns::path)
					.load(connection)?;
			}

			let missing_songs = all_songs
				.into_iter()
				.filter(|ref song_path| {
					        let path = Path::new(&song_path);
					        !path.exists() || self.vfs.real_to_virtual(path).is_err()
					       })
				.collect::<Vec<_>>();

			let connection = self.connection.lock().unwrap();
			let connection = connection.deref();
			diesel::delete(songs::table.filter(songs::columns::path.eq_any(missing_songs)))
				.execute(connection)?;
		}

		{
			let all_directories: Vec<String>;
			{
				let connection = self.connection.lock().unwrap();
				let connection = connection.deref();
				all_directories = directories::table
					.select(directories::columns::path)
					.load(connection)?;
			}

			let missing_directories = all_directories
				.into_iter()
				.filter(|ref directory_path| {
					        let path = Path::new(&directory_path);
					        !path.exists() || self.vfs.real_to_virtual(path).is_err()
					       })
				.collect::<Vec<_>>();

			let connection = self.connection.lock().unwrap();
			let connection = connection.deref();
			diesel::delete(directories::table.filter(directories::columns::path
			                                             .eq_any(missing_directories)))
					.execute(connection)?;
		}

		Ok(())
	}

	fn populate(&self) -> Result<()> {
		let vfs = self.vfs.deref();
		let mount_points = vfs.get_mount_points();
		let mut builder = IndexBuilder::new(&self.connection)?;
		for (_, target) in mount_points {
			self.populate_directory(&mut builder, None, target.as_path())?;
		}
		builder.flush_songs()?;
		builder.flush_directories()?;
		Ok(())
	}

	fn get_artwork(&self, dir: &Path) -> Option<String> {
		let pattern = match self.album_art_pattern {
			Some(ref p) => p,
			_ => return None,
		};

		if let Ok(dir_content) = fs::read_dir(dir) {
			for file in dir_content {
				if let Ok(file) = file {
					if let Some(name_string) = file.file_name().to_str() {
						if pattern.is_match(name_string) {
							return file.path().to_str().map(|p| p.to_owned());
						}
					}
				}
			}
		}

		None
	}

	fn populate_directory(&self,
	                      builder: &mut IndexBuilder,
	                      parent: Option<&Path>,
	                      path: &Path)
	                      -> Result<()> {

		// Find artwork
		let artwork = self.get_artwork(path);

		// Extract path and parent path
		let parent_string = parent.and_then(|p| p.to_str()).map(|s| s.to_owned());
		let path_string = path.to_str().ok_or("Invalid directory path")?;

		// Find date added
		let metadata = fs::metadata(path_string)?;
		let created = metadata
			.created()
			.or(metadata.modified())?
			.duration_since(time::UNIX_EPOCH)?
			.as_secs() as i32;

		let mut directory_album = None;
		let mut directory_year = None;
		let mut directory_artist = None;
		let mut inconsistent_directory_album = false;
		let mut inconsistent_directory_year = false;
		let mut inconsistent_directory_artist = false;

		// Insert content
		if let Ok(dir_content) = fs::read_dir(path) {
			for file in dir_content {
				let file_path = match file {
					Ok(f) => f.path(),
					_ => continue,
				};

				if file_path.is_dir() {
					self.populate_directory(builder, Some(path), file_path.as_path())?;
				} else {
					if let Some(file_path_string) = file_path.to_str() {
						if let Ok(tags) = metadata::read(file_path.as_path()) {
							if tags.year.is_some() {
								inconsistent_directory_year |= directory_year.is_some() &&
								                               directory_year != tags.year;
								directory_year = tags.year;
							}

							if tags.album.is_some() {
								inconsistent_directory_album |= directory_album.is_some() &&
								                                directory_album != tags.album;
								directory_album = tags.album.as_ref().map(|a| a.clone());
							}

							if tags.album_artist.is_some() {
								inconsistent_directory_artist |= directory_artist.is_some() &&
								                                 directory_artist !=
								                                 tags.album_artist;
								directory_artist = tags.album_artist.as_ref().map(|a| a.clone());
							} else if tags.artist.is_some() {
								inconsistent_directory_artist |= directory_artist.is_some() &&
								                                 directory_artist != tags.artist;
								directory_artist = tags.artist.as_ref().map(|a| a.clone());
							}

							let song = NewSong {
								path: file_path_string.to_owned(),
								parent: path_string.to_owned(),
								disc_number: tags.disc_number.map(|n| n as i32),
								track_number: tags.track_number.map(|n| n as i32),
								title: tags.title,
								artist: tags.artist,
								album_artist: tags.album_artist,
								album: tags.album,
								year: tags.year,
								artwork: artwork.as_ref().map(|s| s.to_owned()),
							};

							builder.push_song(song)?;
						}
					}
				}
			}
		}

		// Insert directory
		if inconsistent_directory_year {
			directory_year = None;
		}
		if inconsistent_directory_album {
			directory_album = None;
		}
		if inconsistent_directory_artist {
			directory_artist = None;
		}

		let directory = NewDirectory {
			path: path_string.to_owned(),
			parent: parent_string,
			artwork: artwork,
			album: directory_album,
			artist: directory_artist,
			year: directory_year,
			date_added: created,
		};
		builder.push_directory(directory)?;

		Ok(())
	}

	pub fn update_loop(&self) {
		loop {
			if let Err(e) = self.update_index() {
				println!("Error while updating index: {}", e);
			}
			thread::sleep(time::Duration::from_secs(self.sleep_duration));
		}
	}
}

fn _get_test_db(name: &str) -> DB {
	use vfs::VfsConfig;
	use std::collections::HashMap;

	let mut collection_path = PathBuf::new();
	collection_path.push("test");
	collection_path.push("collection");
	let mut mount_points = HashMap::new();
	mount_points.insert("root".to_owned(), collection_path);

	let vfs = Arc::new(Vfs::new(VfsConfig { mount_points: mount_points }));

	let mut index_config = IndexConfig::new();
	index_config.album_art_pattern = Some(Regex::new(r#"^Folder\.(png|jpg|jpeg)$"#).unwrap());
	index_config.path = PathBuf::new();
	index_config.path.push("test");
	index_config.path.push(name);

	if index_config.path.exists() {
		fs::remove_file(&index_config.path).unwrap();
	}

	DB::new(vfs, &index_config).unwrap()
}

#[test]
fn test_populate() {
	use db::models::*;

	let db = _get_test_db("populate.sqlite");
	let index = db.get_index();
	index.update_index().unwrap();
	index.update_index().unwrap(); // Check that subsequent updates don't run into conflicts

	let connection = index.connection.lock().unwrap();
	let connection = connection.deref();
	let all_directories: Vec<Directory> = directories::table.load(connection).unwrap();
	let all_songs: Vec<Song> = songs::table.load(connection).unwrap();
	assert_eq!(all_directories.len(), 5);
	assert_eq!(all_songs.len(), 12);
}

#[test]
fn test_metadata() {
	use db::models::*;

	let mut target = PathBuf::new();
	target.push("test");
	target.push("collection");
	target.push("Tobokegao");
	target.push("Picnic");

	let mut song_path = target.clone();
	song_path.push("05 - シャーベット (Sherbet).mp3");

	let mut artwork_path = target.clone();
	artwork_path.push("Folder.png");

	let db = _get_test_db("metadata.sqlite");
	let index = db.get_index();
	index.update_index().unwrap();

	let connection = index.connection.lock().unwrap();
	let connection = connection.deref();
	let songs: Vec<Song> = songs::table
		.filter(songs::columns::title.eq("シャーベット (Sherbet)"))
		.load(connection)
		.unwrap();

	assert_eq!(songs.len(), 1);
	let song = &songs[0];
	assert_eq!(song.path, song_path.to_string_lossy().as_ref());
	assert_eq!(song.track_number, Some(5));
	assert_eq!(song.disc_number, None);
	assert_eq!(song.title, Some("シャーベット (Sherbet)".to_owned()));
	assert_eq!(song.artist, Some("Tobokegao".to_owned()));
	assert_eq!(song.album_artist, None);
	assert_eq!(song.album, Some("Picnic".to_owned()));
	assert_eq!(song.year, Some(2016));
	assert_eq!(song.artwork,
	           Some(artwork_path.to_string_lossy().into_owned()));
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::old_io;
use std::old_io::{fs, File};
use std::old_io::fs::PathExtensions;
use std::path::BytesContainer;

use mustache;

use error::{FerrumError, FerrumResult};
use parser;

#[derive(PartialEq, Show)]
pub struct Document {
    data: HashMap<String, String>,
    content: String,
}

impl Document {
    pub fn new(header: HashMap<String, String>, content: &str) -> Document {
        Document { data: header, content: content.to_string() }
    }

    pub fn as_html(&self) -> String {
        let template = mustache::compile_str(self.content.as_slice());

        // Write the template to memory, then retrieve it as a string.
        let mut w = Vec::<u8>::new();
        //let mut w = old_io::MemWriter::new();
        template.render(&mut w, &self.data).is_ok();

        w.container_as_str().unwrap().to_string()
    }

    pub fn render_to_file(&self, file_path: &Path, templates: &HashMap<String, String>) -> FerrumResult<()> {
        let template_path = try!(self.template());
        let template = match templates.get(&template_path.to_string()) {
            Some(template) => template,
            None => return Err(FerrumError::InvalidDocumentError("Template not found".to_string()))
        };

        fs::mkdir_recursive(&file_path.dir_path(), old_io::USER_RWX).is_ok();

        let mut file = File::create(file_path);
        let mut data = HashMap::new();

        data.insert("content", self.as_html());

        for (key, value) in self.data.iter() {
            data.insert(key.as_slice(), value.clone());
        }

        let template = mustache::compile_str(template.as_slice());

        template.render(&mut file, &data).is_ok();

        info!("Created {}", file_path.display());
        Ok(())
    }

    fn template(&self) -> FerrumResult<&str> {
        match self.data.get(&"template".to_string()) {
            Some(v) => Ok(v.as_slice()),
            None => Err(FerrumError::InvalidDocumentError("Missing template".to_string()))
        }
    }
}

pub fn load_documents_from_disk<F>(documents_path: &Path, mut criteria: F) -> FerrumResult<HashMap<Path, Document>>
    where F : FnMut(&Path) -> bool
{
    let mut documents = HashMap::new();

    let mut document_dirs = try!(fs::walk_dir(documents_path));
    for path in document_dirs {
        if !path.is_file() { continue; }
        if !criteria(&path) { continue; }

        // Read the document from the disk.
        let content = try!(File::open(&path).read_to_end());
        let content = String::from_utf8_lossy(content.as_slice());
        let document = match parser::document(content.as_slice()) {
            Ok(document) => document,
            Err(err) => {
                warn!("Failed to read document {}: {}", path.display(), FerrumError::ParserError(err));
                continue;
            }
        };
        let mut relative_path = path.path_relative_from(documents_path).unwrap();
        relative_path.set_extension("");

        documents.insert(relative_path, document);
    }

    Ok(documents)
}
<commit_msg>Fix slicing syntax<commit_after>use std::collections::HashMap;
use std::old_io;
use std::old_io::{fs, File};
use std::old_io::fs::PathExtensions;
use std::path::BytesContainer;

use mustache;

use error::{FerrumError, FerrumResult};
use parser;

#[derive(PartialEq, Show)]
pub struct Document {
    data: HashMap<String, String>,
    content: String,
}

impl Document {
    pub fn new(header: HashMap<String, String>, content: &str) -> Document {
        Document { data: header, content: content.to_string() }
    }

    pub fn as_html(&self) -> String {
        let template = mustache::compile_str(&*self.content);

        // Write the template to memory, then retrieve it as a string.
        let mut w = Vec::<u8>::new();
        //let mut w = old_io::MemWriter::new();
        template.render(&mut w, &self.data).is_ok();

        w.container_as_str().unwrap().to_string()
    }

    pub fn render_to_file(&self, file_path: &Path, templates: &HashMap<String, String>) -> FerrumResult<()> {
        let template_path = try!(self.template());
        let template = match templates.get(&template_path.to_string()) {
            Some(template) => template,
            None => return Err(FerrumError::InvalidDocumentError("Template not found".to_string()))
        };

        fs::mkdir_recursive(&file_path.dir_path(), old_io::USER_RWX).is_ok();

        let mut file = File::create(file_path);
        let mut data = HashMap::new();

        data.insert("content", self.as_html());

        for (key, value) in self.data.iter() {
            data.insert(&key[], value.clone());
        }

        let template = mustache::compile_str(&template[]);

        template.render(&mut file, &data).is_ok();

        info!("Created {}", file_path.display());
        Ok(())
    }

    fn template(&self) -> FerrumResult<&str> {
        match self.data.get(&"template".to_string()) {
            Some(v) => Ok(&v[]),
            None => Err(FerrumError::InvalidDocumentError("Missing template".to_string()))
        }
    }
}

pub fn load_documents_from_disk<F>(documents_path: &Path, mut criteria: F) -> FerrumResult<HashMap<Path, Document>>
    where F : FnMut(&Path) -> bool
{
    let mut documents = HashMap::new();

    let mut document_dirs = try!(fs::walk_dir(documents_path));
    for path in document_dirs {
        if !path.is_file() { continue; }
        if !criteria(&path) { continue; }

        // Read the document from the disk.
        let content = try!(File::open(&path).read_to_end());
        let content = String::from_utf8_lossy(&*content);
        let document = match parser::document(&*content) {
            Ok(document) => document,
            Err(err) => {
                warn!("Failed to read document {}: {}", path.display(), FerrumError::ParserError(err));
                continue;
            }
        };
        let mut relative_path = path.path_relative_from(documents_path).unwrap();
        relative_path.set_extension("");

        documents.insert(relative_path, document);
    }

    Ok(documents)
}
<|endoftext|>


<commit_before>use super::content::Content;
use super::encodings::{self, bytes_to_string, string_to_bytes};
use super::{Dictionary, Object, ObjectId};
use crate::xref::Xref;
use crate::{Error, Result};
use encoding::all::UTF_16BE;
use encoding::types::{DecoderTrap, EncoderTrap, Encoding};
use log::info;
use std::collections::BTreeMap;
use std::io::Write;
use std::str;

/// PDF document.
#[derive(Debug, Clone)]
pub struct Document {
	/// The version of the PDF specification to which the file conforms.
	pub version: String,

	/// The trailer gives the location of the cross-reference table and of certain special objects.
	pub trailer: Dictionary,

	/// The cross-reference table contains locations of the indirect objects.
	pub reference_table: Xref,

	/// The objects that make up the document contained in the file.
	pub objects: BTreeMap<ObjectId, Object>,

	/// Current maximum object id within the document.
	pub max_id: u32,
}

impl Document {
	/// Create new PDF document.
	pub fn new() -> Document {
		Document {
			version: "1.4".to_string(),
			trailer: Dictionary::new(),
			reference_table: Xref::new(0),
			objects: BTreeMap::new(),
			max_id: 0,
		}
	}

	const DEREF_LIMIT : usize = 128;

	/// Follow references if the supplied object is a reference.
	///
	/// Returns a tuple of an optional object id and final object.
	/// The object id will be None if the object was not a
	/// reference. Otherwise, it will be the last object id in the
	/// reference chain.
	pub fn dereference<'a>(&'a self, mut object: &'a Object) -> Result<(Option<ObjectId>, &'a Object)> {
		let mut nb_deref = 0;
		let mut id = None;

		while let Ok(ref_id) = object.as_reference() {
			id = Some(ref_id);
			object = self.objects.get(&ref_id).ok_or(Error::ObjectNotFound)?;

			nb_deref += 1;
			if nb_deref > Self::DEREF_LIMIT {
				return Err(Error::ReferenceLimit);
			}
		}

		Ok((id, object))
	}

	/// Get object by object id, will iteratively dereference a referenced object.
	pub fn get_object(&self, id: ObjectId) -> Result<&Object> {
		let object = self.objects.get(&id).ok_or(Error::ObjectNotFound)?;
		self.dereference(object).map(|(_, object)| object)
	}

	/// Get mutable reference to object by object id, will iteratively dereference a referenced object.
	pub fn get_object_mut(&mut self, id: ObjectId) -> Result<&mut Object> {
		let object = self.objects.get(&id).ok_or(Error::ObjectNotFound)?;
		let (ref_id, _) = self.dereference(object)?;

		Ok(self.objects.get_mut(&ref_id.unwrap_or(id)).unwrap())
	}

	/// Get dictionary object by id.
	pub fn get_dictionary(&self, id: ObjectId) -> Result<&Dictionary> {
		self.get_object(id).and_then(Object::as_dict)
	}

	/// Traverse objects from trailer recursively, return all referenced object IDs.
	pub fn traverse_objects<A: Fn(&mut Object) -> ()>(&mut self, action: A) -> Vec<ObjectId> {
		fn traverse_array<A: Fn(&mut Object) -> ()>(array: &mut Vec<Object>, action: &A, refs: &mut Vec<ObjectId>) {
			for item in array.iter_mut() {
				traverse_object(item, action, refs);
			}
		}
		fn traverse_dictionary<A: Fn(&mut Object) -> ()>(dict: &mut Dictionary, action: &A, refs: &mut Vec<ObjectId>) {
			for (_, v) in dict.iter_mut() {
				traverse_object(v, action, refs);
			}
		}
		fn traverse_object<A: Fn(&mut Object) -> ()>(object: &mut Object, action: &A, refs: &mut Vec<ObjectId>) {
			action(object);
			match *object {
				Object::Array(ref mut array) => traverse_array(array, action, refs),
				Object::Dictionary(ref mut dict) => traverse_dictionary(dict, action, refs),
				Object::Stream(ref mut stream) => traverse_dictionary(&mut stream.dict, action, refs),
				Object::Reference(id) => {
					if !refs.contains(&id) {
						refs.push(id);
					}
				}
				_ => {}
			}
		}
		let mut refs = vec![];
		traverse_dictionary(&mut self.trailer, &action, &mut refs);
		let mut index = 0;
		while index < refs.len() {
			if let Some(object) = self.objects.get_mut(&refs[index]) {
				traverse_object(object, &action, &mut refs);
			}
			index += 1;
		}
		refs
	}

	/// Get catalog dictionary.
	pub fn catalog(&self) -> Result<&Dictionary> {
		self.trailer.get(b"Root").and_then(Object::as_reference).and_then(|id| self.get_dictionary(id))
	}

	/// Get page numbers and corresponding object ids.
	pub fn get_pages(&self) -> BTreeMap<u32, ObjectId> {
		self.page_iter().enumerate().map(|(i, p)| ((i+1) as u32, p)).collect()
	}

	pub fn page_iter(&self) -> impl Iterator<Item=ObjectId> + '_ {
		PageTreeIter::new(self)
	}

	/// Get content stream object ids of a page.
	pub fn get_page_contents(&self, page_id: ObjectId) -> Vec<ObjectId> {
		let mut streams = vec![];
		if let Ok(page) = self.get_dictionary(page_id) {
			if let Ok(contents) = page.get(b"Contents") {
				match *contents {
					Object::Reference(ref id) => {
						streams.push(*id);
					}
					Object::Array(ref arr) => {
						for content in arr {
							if let Ok(id) = content.as_reference() { streams.push(id) }
						}
					}
					_ => {}
				}
			}
		}
		streams
	}

	/// Get content of a page.
	pub fn get_page_content(&self, page_id: ObjectId) -> Result<Vec<u8>> {
		let mut content = Vec::new();
		let content_streams = self.get_page_contents(page_id);
		for object_id in content_streams {
			if let Ok(content_stream) = self.get_object(object_id).and_then(Object::as_stream) {
				match content_stream.decompressed_content() {
					Ok(data) => content.write_all(&data)?,
					Err(_) => content.write_all(&content_stream.content)?,
				};
			}
		}
		Ok(content)
	}

	/// Get decoded page content;
	pub fn get_and_decode_page_content(&self, page_id: ObjectId) -> Result<Content> {
		let content_data = self.get_page_content(page_id)?;
		Content::decode(&content_data)
	}

	/// Get resources used by a page.
	pub fn get_page_resources(&self, page_id: ObjectId) -> (Option<&Dictionary>, Vec<ObjectId>) {
		fn collect_resources(page_node: &Dictionary, resource_ids: &mut Vec<ObjectId>, doc: &Document) {
			if let Ok(resources_id) = page_node.get(b"Resources").and_then(Object::as_reference) {
				resource_ids.push(resources_id);
			}
			if let Ok(page_tree) = page_node.get(b"Parent").and_then(Object::as_reference).and_then(|id| doc.get_dictionary(id)) {
				collect_resources(page_tree, resource_ids, doc);
			}
		};

		let mut resource_dict = None;
		let mut resource_ids = Vec::new();
		if let Ok(page) = self.get_dictionary(page_id) {
			resource_dict = page.get(b"Resources").and_then(Object::as_dict).ok();
			collect_resources(page, &mut resource_ids, self);
		}
		(resource_dict, resource_ids)
	}

	/// Get fonts used by a page.
	pub fn get_page_fonts(&self, page_id: ObjectId) -> BTreeMap<Vec<u8>, &Dictionary> {
		fn collect_fonts_from_resources<'a>(resources: &'a Dictionary, fonts: &mut BTreeMap<Vec<u8>, &'a Dictionary>, doc: &'a Document) {
			if let Ok(font_dict) = resources.get(b"Font").and_then(Object::as_dict) {
				for (name, value) in font_dict.iter() {
					let font = match *value {
						Object::Reference(id) => doc.get_dictionary(id).ok(),
						Object::Dictionary(ref dict) => Some(dict),
						_ => None,
					};
					if !fonts.contains_key(name) {
						font.map(|font| fonts.insert(name.clone(), font));
					}
				}
			}
		};

		let mut fonts = BTreeMap::new();
		let (resource_dict, resource_ids) = self.get_page_resources(page_id);
		if let Some(resources) = resource_dict {
			collect_fonts_from_resources(resources, &mut fonts, self);
		}
		for resource_id in resource_ids {
			if let Ok(resources) = self.get_dictionary(resource_id) {
				collect_fonts_from_resources(resources, &mut fonts, self);
			}
		}
		fonts
	}

	pub fn decode_text(encoding: Option<&str>, bytes: &[u8]) -> String {
		if let Some(encoding) = encoding {
			info!("{}", encoding);
			match encoding {
				"StandardEncoding" => bytes_to_string(encodings::STANDARD_ENCODING, bytes),
				"MacRomanEncoding" => bytes_to_string(encodings::MAC_ROMAN_ENCODING, bytes),
				"MacExpertEncoding" => bytes_to_string(encodings::MAC_EXPERT_ENCODING, bytes),
				"WinAnsiEncoding" => bytes_to_string(encodings::WIN_ANSI_ENCODING, bytes),
				"UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.decode(bytes, DecoderTrap::Ignore).unwrap(),
				"Identity-H" => "?Identity-H Unimplemented?".to_string(), // Unimplemented
				_ => String::from_utf8_lossy(bytes).to_string(),
			}
		} else {
			bytes_to_string(encodings::STANDARD_ENCODING, bytes)
		}
	}

	pub fn encode_text(encoding: Option<&str>, text: &str) -> Vec<u8> {
		if let Some(encoding) = encoding {
			match encoding {
				"StandardEncoding" => string_to_bytes(encodings::STANDARD_ENCODING, text),
				"MacRomanEncoding" => string_to_bytes(encodings::MAC_ROMAN_ENCODING, text),
				"MacExpertEncoding" => string_to_bytes(encodings::MAC_EXPERT_ENCODING, text),
				"WinAnsiEncoding" => string_to_bytes(encodings::WIN_ANSI_ENCODING, text),
				"UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.encode(text, EncoderTrap::Ignore).unwrap(),
				"Identity-H" => vec![], // Unimplemented
				_ => text.as_bytes().to_vec(),
			}
		} else {
			string_to_bytes(encodings::STANDARD_ENCODING, text)
		}
	}
}

impl Default for Document {
	fn default() -> Self {
		Self::new()
	}
}

struct PageTreeIter<'a> {
	doc: &'a Document,
	stack: Vec<&'a [Object]>,
	kids: Option<&'a [Object]>,
}

impl <'a> PageTreeIter<'a> {
	const PAGE_TREE_DEPTH_LIMIT: usize = 256;

	fn new(doc: &'a Document) -> Self {
		if let Ok(page_tree_id) = doc.catalog().and_then(|cat| cat.get(b"Pages")).and_then(Object::as_reference) {
			Self { doc, kids: Self::kids(doc, page_tree_id), stack: Vec::with_capacity(32) }
		} else {
			Self { doc, kids: None, stack: Vec::new() }
		}
	}

	fn kids(doc: &Document, page_tree_id: ObjectId) -> Option<&[Object]> {
		doc.get_dictionary(page_tree_id).and_then(|page_tree| page_tree.get(b"Kids")).and_then(Object::as_array).map(|k| k.as_slice()).ok()
	}
}

impl Iterator for PageTreeIter<'_> {
	type Item = ObjectId;

	fn next(&mut self) -> Option<Self::Item> {
		loop {
			while let Some((kid, new_kids)) = self.kids.and_then(|k| k.split_first()) {
				self.kids = Some(new_kids);

				if let Ok(kid_id) = kid.as_reference() {
					if let Ok(type_name) = self.doc.get_dictionary(kid_id).and_then(Dictionary::type_name) {
						match type_name {
							"Page" => {
								return Some(kid_id);
							},
							"Pages" => {
								if self.stack.len() < Self::PAGE_TREE_DEPTH_LIMIT {
									let kids = self.kids.unwrap();
									if !kids.is_empty() {
										self.stack.push(kids);
									}
									self.kids = Self::kids(self.doc, kid_id);
								}
							},
							_ => {}
						}
					}
				}
			}

			// Current level exhausted, try to pop.
			if let kids @ Some(_) = self.stack.pop() {
				self.kids = kids;
			} else {
				return None;
			}
		}
	}

	fn size_hint(&self) -> (usize, Option<usize>) {
		let stack_len: usize = self.stack.iter().map(|s| s.len()).sum();

		(self.kids.map(|k| k.len()).unwrap_or(0) + stack_len, None)
	}
}

impl std::iter::FusedIterator for PageTreeIter<'_> {}
<commit_msg>Compute an accurate iterator size when the page tree is sane.<commit_after>use super::content::Content;
use super::encodings::{self, bytes_to_string, string_to_bytes};
use super::{Dictionary, Object, ObjectId};
use crate::xref::Xref;
use crate::{Error, Result};
use encoding::all::UTF_16BE;
use encoding::types::{DecoderTrap, EncoderTrap, Encoding};
use log::info;
use std::collections::BTreeMap;
use std::cmp::max;
use std::io::Write;
use std::str;

/// PDF document.
#[derive(Debug, Clone)]
pub struct Document {
	/// The version of the PDF specification to which the file conforms.
	pub version: String,

	/// The trailer gives the location of the cross-reference table and of certain special objects.
	pub trailer: Dictionary,

	/// The cross-reference table contains locations of the indirect objects.
	pub reference_table: Xref,

	/// The objects that make up the document contained in the file.
	pub objects: BTreeMap<ObjectId, Object>,

	/// Current maximum object id within the document.
	pub max_id: u32,
}

impl Document {
	/// Create new PDF document.
	pub fn new() -> Document {
		Document {
			version: "1.4".to_string(),
			trailer: Dictionary::new(),
			reference_table: Xref::new(0),
			objects: BTreeMap::new(),
			max_id: 0,
		}
	}

	const DEREF_LIMIT : usize = 128;

	/// Follow references if the supplied object is a reference.
	///
	/// Returns a tuple of an optional object id and final object.
	/// The object id will be None if the object was not a
	/// reference. Otherwise, it will be the last object id in the
	/// reference chain.
	pub fn dereference<'a>(&'a self, mut object: &'a Object) -> Result<(Option<ObjectId>, &'a Object)> {
		let mut nb_deref = 0;
		let mut id = None;

		while let Ok(ref_id) = object.as_reference() {
			id = Some(ref_id);
			object = self.objects.get(&ref_id).ok_or(Error::ObjectNotFound)?;

			nb_deref += 1;
			if nb_deref > Self::DEREF_LIMIT {
				return Err(Error::ReferenceLimit);
			}
		}

		Ok((id, object))
	}

	/// Get object by object id, will iteratively dereference a referenced object.
	pub fn get_object(&self, id: ObjectId) -> Result<&Object> {
		let object = self.objects.get(&id).ok_or(Error::ObjectNotFound)?;
		self.dereference(object).map(|(_, object)| object)
	}

	/// Get mutable reference to object by object id, will iteratively dereference a referenced object.
	pub fn get_object_mut(&mut self, id: ObjectId) -> Result<&mut Object> {
		let object = self.objects.get(&id).ok_or(Error::ObjectNotFound)?;
		let (ref_id, _) = self.dereference(object)?;

		Ok(self.objects.get_mut(&ref_id.unwrap_or(id)).unwrap())
	}

	/// Get dictionary object by id.
	pub fn get_dictionary(&self, id: ObjectId) -> Result<&Dictionary> {
		self.get_object(id).and_then(Object::as_dict)
	}

	/// Traverse objects from trailer recursively, return all referenced object IDs.
	pub fn traverse_objects<A: Fn(&mut Object) -> ()>(&mut self, action: A) -> Vec<ObjectId> {
		fn traverse_array<A: Fn(&mut Object) -> ()>(array: &mut Vec<Object>, action: &A, refs: &mut Vec<ObjectId>) {
			for item in array.iter_mut() {
				traverse_object(item, action, refs);
			}
		}
		fn traverse_dictionary<A: Fn(&mut Object) -> ()>(dict: &mut Dictionary, action: &A, refs: &mut Vec<ObjectId>) {
			for (_, v) in dict.iter_mut() {
				traverse_object(v, action, refs);
			}
		}
		fn traverse_object<A: Fn(&mut Object) -> ()>(object: &mut Object, action: &A, refs: &mut Vec<ObjectId>) {
			action(object);
			match *object {
				Object::Array(ref mut array) => traverse_array(array, action, refs),
				Object::Dictionary(ref mut dict) => traverse_dictionary(dict, action, refs),
				Object::Stream(ref mut stream) => traverse_dictionary(&mut stream.dict, action, refs),
				Object::Reference(id) => {
					if !refs.contains(&id) {
						refs.push(id);
					}
				}
				_ => {}
			}
		}
		let mut refs = vec![];
		traverse_dictionary(&mut self.trailer, &action, &mut refs);
		let mut index = 0;
		while index < refs.len() {
			if let Some(object) = self.objects.get_mut(&refs[index]) {
				traverse_object(object, &action, &mut refs);
			}
			index += 1;
		}
		refs
	}

	/// Get catalog dictionary.
	pub fn catalog(&self) -> Result<&Dictionary> {
		self.trailer.get(b"Root").and_then(Object::as_reference).and_then(|id| self.get_dictionary(id))
	}

	/// Get page numbers and corresponding object ids.
	pub fn get_pages(&self) -> BTreeMap<u32, ObjectId> {
		self.page_iter().enumerate().map(|(i, p)| ((i+1) as u32, p)).collect()
	}

	pub fn page_iter(&self) -> impl Iterator<Item=ObjectId> + '_ {
		PageTreeIter::new(self)
	}

	/// Get content stream object ids of a page.
	pub fn get_page_contents(&self, page_id: ObjectId) -> Vec<ObjectId> {
		let mut streams = vec![];
		if let Ok(page) = self.get_dictionary(page_id) {
			if let Ok(contents) = page.get(b"Contents") {
				match *contents {
					Object::Reference(ref id) => {
						streams.push(*id);
					}
					Object::Array(ref arr) => {
						for content in arr {
							if let Ok(id) = content.as_reference() { streams.push(id) }
						}
					}
					_ => {}
				}
			}
		}
		streams
	}

	/// Get content of a page.
	pub fn get_page_content(&self, page_id: ObjectId) -> Result<Vec<u8>> {
		let mut content = Vec::new();
		let content_streams = self.get_page_contents(page_id);
		for object_id in content_streams {
			if let Ok(content_stream) = self.get_object(object_id).and_then(Object::as_stream) {
				match content_stream.decompressed_content() {
					Ok(data) => content.write_all(&data)?,
					Err(_) => content.write_all(&content_stream.content)?,
				};
			}
		}
		Ok(content)
	}

	/// Get decoded page content;
	pub fn get_and_decode_page_content(&self, page_id: ObjectId) -> Result<Content> {
		let content_data = self.get_page_content(page_id)?;
		Content::decode(&content_data)
	}

	/// Get resources used by a page.
	pub fn get_page_resources(&self, page_id: ObjectId) -> (Option<&Dictionary>, Vec<ObjectId>) {
		fn collect_resources(page_node: &Dictionary, resource_ids: &mut Vec<ObjectId>, doc: &Document) {
			if let Ok(resources_id) = page_node.get(b"Resources").and_then(Object::as_reference) {
				resource_ids.push(resources_id);
			}
			if let Ok(page_tree) = page_node.get(b"Parent").and_then(Object::as_reference).and_then(|id| doc.get_dictionary(id)) {
				collect_resources(page_tree, resource_ids, doc);
			}
		};

		let mut resource_dict = None;
		let mut resource_ids = Vec::new();
		if let Ok(page) = self.get_dictionary(page_id) {
			resource_dict = page.get(b"Resources").and_then(Object::as_dict).ok();
			collect_resources(page, &mut resource_ids, self);
		}
		(resource_dict, resource_ids)
	}

	/// Get fonts used by a page.
	pub fn get_page_fonts(&self, page_id: ObjectId) -> BTreeMap<Vec<u8>, &Dictionary> {
		fn collect_fonts_from_resources<'a>(resources: &'a Dictionary, fonts: &mut BTreeMap<Vec<u8>, &'a Dictionary>, doc: &'a Document) {
			if let Ok(font_dict) = resources.get(b"Font").and_then(Object::as_dict) {
				for (name, value) in font_dict.iter() {
					let font = match *value {
						Object::Reference(id) => doc.get_dictionary(id).ok(),
						Object::Dictionary(ref dict) => Some(dict),
						_ => None,
					};
					if !fonts.contains_key(name) {
						font.map(|font| fonts.insert(name.clone(), font));
					}
				}
			}
		};

		let mut fonts = BTreeMap::new();
		let (resource_dict, resource_ids) = self.get_page_resources(page_id);
		if let Some(resources) = resource_dict {
			collect_fonts_from_resources(resources, &mut fonts, self);
		}
		for resource_id in resource_ids {
			if let Ok(resources) = self.get_dictionary(resource_id) {
				collect_fonts_from_resources(resources, &mut fonts, self);
			}
		}
		fonts
	}

	pub fn decode_text(encoding: Option<&str>, bytes: &[u8]) -> String {
		if let Some(encoding) = encoding {
			info!("{}", encoding);
			match encoding {
				"StandardEncoding" => bytes_to_string(encodings::STANDARD_ENCODING, bytes),
				"MacRomanEncoding" => bytes_to_string(encodings::MAC_ROMAN_ENCODING, bytes),
				"MacExpertEncoding" => bytes_to_string(encodings::MAC_EXPERT_ENCODING, bytes),
				"WinAnsiEncoding" => bytes_to_string(encodings::WIN_ANSI_ENCODING, bytes),
				"UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.decode(bytes, DecoderTrap::Ignore).unwrap(),
				"Identity-H" => "?Identity-H Unimplemented?".to_string(), // Unimplemented
				_ => String::from_utf8_lossy(bytes).to_string(),
			}
		} else {
			bytes_to_string(encodings::STANDARD_ENCODING, bytes)
		}
	}

	pub fn encode_text(encoding: Option<&str>, text: &str) -> Vec<u8> {
		if let Some(encoding) = encoding {
			match encoding {
				"StandardEncoding" => string_to_bytes(encodings::STANDARD_ENCODING, text),
				"MacRomanEncoding" => string_to_bytes(encodings::MAC_ROMAN_ENCODING, text),
				"MacExpertEncoding" => string_to_bytes(encodings::MAC_EXPERT_ENCODING, text),
				"WinAnsiEncoding" => string_to_bytes(encodings::WIN_ANSI_ENCODING, text),
				"UniGB-UCS2-H" | "UniGB−UTF16−H" => UTF_16BE.encode(text, EncoderTrap::Ignore).unwrap(),
				"Identity-H" => vec![], // Unimplemented
				_ => text.as_bytes().to_vec(),
			}
		} else {
			string_to_bytes(encodings::STANDARD_ENCODING, text)
		}
	}
}

impl Default for Document {
	fn default() -> Self {
		Self::new()
	}
}

struct PageTreeIter<'a> {
	doc: &'a Document,
	stack: Vec<&'a [Object]>,
	kids: Option<&'a [Object]>,
}

impl <'a> PageTreeIter<'a> {
	const PAGE_TREE_DEPTH_LIMIT: usize = 256;

	fn new(doc: &'a Document) -> Self {
		if let Ok(page_tree_id) = doc.catalog().and_then(|cat| cat.get(b"Pages")).and_then(Object::as_reference) {
			Self { doc, kids: Self::kids(doc, page_tree_id), stack: Vec::with_capacity(32) }
		} else {
			Self { doc, kids: None, stack: Vec::new() }
		}
	}

	fn kids(doc: &Document, page_tree_id: ObjectId) -> Option<&[Object]> {
		doc.get_dictionary(page_tree_id).and_then(|page_tree| page_tree.get(b"Kids")).and_then(Object::as_array).map(|k| k.as_slice()).ok()
	}
}

impl Iterator for PageTreeIter<'_> {
	type Item = ObjectId;

	fn next(&mut self) -> Option<Self::Item> {
		loop {
			while let Some((kid, new_kids)) = self.kids.and_then(|k| k.split_first()) {
				self.kids = Some(new_kids);

				if let Ok(kid_id) = kid.as_reference() {
					if let Ok(type_name) = self.doc.get_dictionary(kid_id).and_then(Dictionary::type_name) {
						match type_name {
							"Page" => {
								return Some(kid_id);
							},
							"Pages" => {
								if self.stack.len() < Self::PAGE_TREE_DEPTH_LIMIT {
									let kids = self.kids.unwrap();
									if !kids.is_empty() {
										self.stack.push(kids);
									}
									self.kids = Self::kids(self.doc, kid_id);
								}
							},
							_ => {}
						}
					}
				}
			}

			// Current level exhausted, try to pop.
			if let kids @ Some(_) = self.stack.pop() {
				self.kids = kids;
			} else {
				return None;
			}
		}
	}

	fn size_hint(&self) -> (usize, Option<usize>) {
		let kids = self.kids.unwrap_or(&[]);

		let nb_pages: usize = kids.iter().chain(self.stack.iter().flat_map(|k| k.iter())).map(|kid| {
			if let Ok(dict) = kid.as_reference().and_then(|id| self.doc.get_dictionary(id)) {
				if let Ok("Pages") = dict.type_name() {
					let count = dict.get_deref(b"Count", self.doc).and_then(Object::as_i64).unwrap_or(0);
					// Don't let page count go backwards in case of an invalid document.
					max(0, count) as usize
				} else {
					1
				}
			} else {
				1
			}
		}).sum();

		(nb_pages, Some(nb_pages))
	}
}

impl std::iter::FusedIterator for PageTreeIter<'_> {}
<|endoftext|>


<commit_before>extern crate std;
extern crate cql;

fn main() {
    println!("Connecting ...!");
    let mut client = match cql::connect("127.0.0.1", 9042, None) {
        Ok(c) => c,
        Err(e) => fail!(format!("Failed to connect: {}", e.desc))
    };
    println!("Connected with CQL binary version v{}", client.version);

    let mut q = "create keyspace rust with replication = {'class': 'SimpleStrategy', 'replication_factor':1}";
    println!("Query: {}", q);
    let mut res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "create table rust.test (id text primary key, f32 float, f64 double, i32 int, i64 bigint, b boolean, ip inet)";
    println!("Query: {}", q);
    res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "insert into rust.test (id, f32, f64, i32, i64, b, ip) values ('asdf', 1.2345, 3.14159, 47, 59, true, '127.0.0.1')";
    println!("Query: {}", q);
    res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "select * from rust.test";
    println!("Query: {}", q);
    res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "insert into rust.test (id, f32) values (?, ?)";
    println!("Create prepared: {}", q);
    let res_id = client.prepared_statement(q);
    println!("Result prepared: {} \n", res);

    if res_id.is_err() {
        fail!("Error in creating prepared statement")
    }

    let ps_id = res_id.unwrap();
    println!("Execute prepared: {}", ps_id);
    let params: Vec<cql::CqlValue> = vec![cql::CqlVarchar(Some("ttrwe".into_maybe_owned())), cql::CqlFloat(Some(15.1617))];
    res = client.exec_prepared(ps_id, params, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "select * from rust.test";
    println!("Query: {}", q);
    res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "create table rust.test2 (id text primary key, l list<int>, m map<int, text>, s set<float>)";
    println!("Query: {}", q);
    res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "insert into rust.test2 (id, l, m, s) values ('asdf', [1,2,3,4,5], {0: 'aa', 1: 'bbb', 3: 'cccc'}, {1.234, 2.345, 3.456, 4.567})";
    println!("Query: {}", q);
    res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "select * from rust.test2";
    println!("Query: {}", q);
    res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);
}
<commit_msg>Use tag to create prepared statement. Use [] instead of Vec<commit_after>extern crate std;
extern crate cql;

fn main() {
    println!("Connecting ...!");
    let mut client = match cql::connect("127.0.0.1", 9042, None) {
        Ok(c) => c,
        Err(e) => fail!(format!("Failed to connect: {}", e.desc))
    };
    println!("Connected with CQL binary version v{}", client.version);

    let mut q = "create keyspace rust with replication = {'class': 'SimpleStrategy', 'replication_factor':1}";
    println!("Query: {}", q);
    let mut res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "create table rust.test (id text primary key, f32 float, f64 double, i32 int, i64 bigint, b boolean, ip inet)";
    println!("Query: {}", q);
    res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "insert into rust.test (id, f32, f64, i32, i64, b, ip) values ('asdf', 1.2345, 3.14159, 47, 59, true, '127.0.0.1')";
    println!("Query: {}", q);
    res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "select * from rust.test";
    println!("Query: {}", q);
    res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "insert into rust.test (id, f32) values (?, ?)";
    println!("Create prepared: {}", q);
    let res_id = client.prepared_statement(q, "test");
    println!("Result prepared: {} \n", res);

    if res_id.is_err() {
        fail!("Error in creating prepared statement")
    }

    println!("Execute prepared");
    let params: &[cql::CqlValue] = &[cql::CqlVarchar(Some("ttrwe".into_maybe_owned())), cql::CqlFloat(Some(15.1617))];
    res = client.exec_prepared("test", params, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "select * from rust.test";
    println!("Query: {}", q);
    res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "create table rust.test2 (id text primary key, l list<int>, m map<int, text>, s set<float>)";
    println!("Query: {}", q);
    res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "insert into rust.test2 (id, l, m, s) values ('asdf', [1,2,3,4,5], {0: 'aa', 1: 'bbb', 3: 'cccc'}, {1.234, 2.345, 3.456, 4.567})";
    println!("Query: {}", q);
    res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);

    q = "select * from rust.test2";
    println!("Query: {}", q);
    res = client.exec_query(q, cql::Consistency::One);
    println!("Result: {} \n", res);
}
<|endoftext|>


<commit_before>//! # The Redox OS Kernel, version 2
//!
//! The Redox OS Kernel is a hybrid kernel that supports X86 systems and
//! provides Unix-like syscalls for primarily Rust applications
//!
//! ## Syscalls
//! Syscalls in Redox are often handled by userspace `schemes`.
//! The essential syscalls in Redox are as follows:
//!
//! ### Open
//! `open(path: &str, flags: usize) -> Result<file_descriptor: usize>`
//!
//! Open a file, providing a path as a `&str` and flags, defined elsewhere.
//!
//! Returns a number, known as a file descriptor, that is passed to other syscalls
//!
//! ### Close
//! `close(file_descriptor: usize) -> Result<()>`
//!
//! Close a file descriptor, providing the file descriptor from `open`
//!
//! Returns an error, `EBADF`, if the file descriptor was not found.
//!
//! This potential error is often ignored by userspace
//!
//! ### Duplicate
//! `dup(file_descriptor: usize) -> Result<file_descriptor: usize>`
//!
//! Duplicate a file descriptor, providing the file descriptor from `open`
//!
//! Returns a new file descriptor, or an error
//!
//! ### Read
//! `read(file_descriptor: usize, buffer: &mut [u8]) -> Result<count: usize>`
//!
//! Read from a file descriptor, providing the file descriptor from `open` and a mutable buffer
//!
//! Returns the number of bytes actually read, or an error
//!
//! ### Write
//! `write(file_descriptor: usize, buffer: &[u8]) -> Result<count: usize>`
//!
//! Write to a file descriptor, providing the file descriptor from `open` and a const buffer
//!
//! Returns the number of bytes actually written, or an error
//!
//! ### Stat
//! `fstat(file_descriptor: usize, stat: &mut Stat) -> Result<()>`
//!
//! Get information from a file descriptor, providing the file descriptor from `open`
//! and a mutable Stat struct, defined elsewhere.
//!
//! Returns an error if the operation failed
//!
//! ### Path
//! `fpath(file_descriptor: usize, buffer: &mut [u8]) -> Result<count: usize>`
//!
//! Read the path of a file descriptor, providing the file descriptor from `open` and
//! a mutable buffer.
//!
//! Returns the number of bytes actually read, or an error
//!
//! The buffer should be 4096 bytes, to ensure that the entire path will fit.
//! An error will be returned, `ENOBUFS`, if the buffer is not long enough for the name.
//! In this case, it is recommended to add one page, 4096 bytes, to the buffer and retry.

#![feature(alloc)]
#![feature(arc_counts)]
#![feature(asm)]
#![feature(collections)]
#![feature(const_fn)]
#![feature(core_intrinsics)]
#![feature(drop_types_in_const)]
#![feature(heap_api)]
#![feature(integer_atomics)]
#![feature(never_type)]
#![feature(thread_local)]
#![no_std]

use arch::interrupt;

/// Architecture specific items (test)
#[cfg(test)]
#[macro_use]
extern crate arch_test as arch;

/// Architecture specific items (ARM)
#[cfg(all(not(test), target_arch = "arm"))]
#[macro_use]
extern crate arch_arm as arch;

/// Architecture specific items (x86_64)
#[cfg(all(not(test), target_arch = "x86_64"))]
#[macro_use]
extern crate arch_x86_64 as arch;

extern crate alloc;
#[macro_use]
extern crate collections;

#[macro_use]
extern crate bitflags;
extern crate goblin;
extern crate spin;

use core::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT, Ordering};

/// Context management
pub mod context;

/// ELF file parsing
#[cfg(all(not(test), target_arch = "x86_64"))]
pub mod elf;

/// Schemes, filesystem handlers
pub mod scheme;

/// Synchronization primitives
pub mod sync;

/// Syscall handlers
pub mod syscall;

/// Tests
#[cfg(test)]
pub mod tests;

#[thread_local]
static CPU_ID: AtomicUsize = ATOMIC_USIZE_INIT;

#[inline(always)]
pub fn cpu_id() -> usize {
    CPU_ID.load(Ordering::Relaxed)
}

static CPU_COUNT : AtomicUsize = ATOMIC_USIZE_INIT;

#[inline(always)]
pub fn cpu_count() -> usize {
    CPU_COUNT.load(Ordering::Relaxed)
}

pub extern fn userspace_init() {
    assert_eq!(syscall::chdir(b"initfs:bin"), Ok(0));

    assert_eq!(syscall::open(b"debug:", 0), Ok(0));
    assert_eq!(syscall::open(b"debug:", 0), Ok(1));
    assert_eq!(syscall::open(b"debug:", 0), Ok(2));

    syscall::exec(b"initfs:bin/init", &[]).expect("failed to execute initfs:init");

    panic!("initfs:init returned")
}

#[no_mangle]
pub extern fn kmain(cpus: usize) {
    CPU_ID.store(0, Ordering::SeqCst);
    CPU_COUNT.store(cpus, Ordering::SeqCst);

    context::init();

    let pid = syscall::getpid();
    println!("BSP: {:?} {}", pid, cpus);

    match context::contexts_mut().spawn(userspace_init) {
        Ok(context_lock) => {
            let mut context = context_lock.write();
            context.status = context::Status::Runnable;
        },
        Err(err) => {
            panic!("failed to spawn userspace_init: {:?}", err);
        }
    }

    loop {
        unsafe {
            interrupt::disable();
            if context::switch() {
                interrupt::enable_and_nop();
            } else {
                interrupt::enable_and_halt();
            }
        }
    }
}

#[no_mangle]
pub extern fn kmain_ap(id: usize) {
    CPU_ID.store(id, Ordering::SeqCst);

    context::init();

    let pid = syscall::getpid();
    println!("AP {}: {:?}", id, pid);

    loop {
        unsafe { interrupt::enable_and_halt() }
    }
}
<commit_msg>Documenting use of `enable_and_halt()`. (#21)<commit_after>//! # The Redox OS Kernel, version 2
//!
//! The Redox OS Kernel is a hybrid kernel that supports X86 systems and
//! provides Unix-like syscalls for primarily Rust applications
//!
//! ## Syscalls
//! Syscalls in Redox are often handled by userspace `schemes`.
//! The essential syscalls in Redox are as follows:
//!
//! ### Open
//! `open(path: &str, flags: usize) -> Result<file_descriptor: usize>`
//!
//! Open a file, providing a path as a `&str` and flags, defined elsewhere.
//!
//! Returns a number, known as a file descriptor, that is passed to other syscalls
//!
//! ### Close
//! `close(file_descriptor: usize) -> Result<()>`
//!
//! Close a file descriptor, providing the file descriptor from `open`
//!
//! Returns an error, `EBADF`, if the file descriptor was not found.
//!
//! This potential error is often ignored by userspace
//!
//! ### Duplicate
//! `dup(file_descriptor: usize) -> Result<file_descriptor: usize>`
//!
//! Duplicate a file descriptor, providing the file descriptor from `open`
//!
//! Returns a new file descriptor, or an error
//!
//! ### Read
//! `read(file_descriptor: usize, buffer: &mut [u8]) -> Result<count: usize>`
//!
//! Read from a file descriptor, providing the file descriptor from `open` and a mutable buffer
//!
//! Returns the number of bytes actually read, or an error
//!
//! ### Write
//! `write(file_descriptor: usize, buffer: &[u8]) -> Result<count: usize>`
//!
//! Write to a file descriptor, providing the file descriptor from `open` and a const buffer
//!
//! Returns the number of bytes actually written, or an error
//!
//! ### Stat
//! `fstat(file_descriptor: usize, stat: &mut Stat) -> Result<()>`
//!
//! Get information from a file descriptor, providing the file descriptor from `open`
//! and a mutable Stat struct, defined elsewhere.
//!
//! Returns an error if the operation failed
//!
//! ### Path
//! `fpath(file_descriptor: usize, buffer: &mut [u8]) -> Result<count: usize>`
//!
//! Read the path of a file descriptor, providing the file descriptor from `open` and
//! a mutable buffer.
//!
//! Returns the number of bytes actually read, or an error
//!
//! The buffer should be 4096 bytes, to ensure that the entire path will fit.
//! An error will be returned, `ENOBUFS`, if the buffer is not long enough for the name.
//! In this case, it is recommended to add one page, 4096 bytes, to the buffer and retry.

#![feature(alloc)]
#![feature(arc_counts)]
#![feature(asm)]
#![feature(collections)]
#![feature(const_fn)]
#![feature(core_intrinsics)]
#![feature(drop_types_in_const)]
#![feature(heap_api)]
#![feature(integer_atomics)]
#![feature(never_type)]
#![feature(thread_local)]
#![no_std]

use arch::interrupt;

/// Architecture specific items (test)
#[cfg(test)]
#[macro_use]
extern crate arch_test as arch;

/// Architecture specific items (ARM)
#[cfg(all(not(test), target_arch = "arm"))]
#[macro_use]
extern crate arch_arm as arch;

/// Architecture specific items (x86_64)
#[cfg(all(not(test), target_arch = "x86_64"))]
#[macro_use]
extern crate arch_x86_64 as arch;

extern crate alloc;
#[macro_use]
extern crate collections;

#[macro_use]
extern crate bitflags;
extern crate goblin;
extern crate spin;

use core::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT, Ordering};

/// Context management
pub mod context;

/// ELF file parsing
#[cfg(all(not(test), target_arch = "x86_64"))]
pub mod elf;

/// Schemes, filesystem handlers
pub mod scheme;

/// Synchronization primitives
pub mod sync;

/// Syscall handlers
pub mod syscall;

/// Tests
#[cfg(test)]
pub mod tests;

#[thread_local]
static CPU_ID: AtomicUsize = ATOMIC_USIZE_INIT;

#[inline(always)]
pub fn cpu_id() -> usize {
    CPU_ID.load(Ordering::Relaxed)
}

static CPU_COUNT : AtomicUsize = ATOMIC_USIZE_INIT;

#[inline(always)]
pub fn cpu_count() -> usize {
    CPU_COUNT.load(Ordering::Relaxed)
}

pub extern fn userspace_init() {
    assert_eq!(syscall::chdir(b"initfs:bin"), Ok(0));

    assert_eq!(syscall::open(b"debug:", 0), Ok(0));
    assert_eq!(syscall::open(b"debug:", 0), Ok(1));
    assert_eq!(syscall::open(b"debug:", 0), Ok(2));

    syscall::exec(b"initfs:bin/init", &[]).expect("failed to execute initfs:init");

    panic!("initfs:init returned")
}

#[no_mangle]
pub extern fn kmain(cpus: usize) {
    CPU_ID.store(0, Ordering::SeqCst);
    CPU_COUNT.store(cpus, Ordering::SeqCst);

    context::init();

    let pid = syscall::getpid();
    println!("BSP: {:?} {}", pid, cpus);

    match context::contexts_mut().spawn(userspace_init) {
        Ok(context_lock) => {
            let mut context = context_lock.write();
            context.status = context::Status::Runnable;
        },
        Err(err) => {
            panic!("failed to spawn userspace_init: {:?}", err);
        }
    }

    loop {
        unsafe {
            interrupt::disable();
            if context::switch() {
                interrupt::enable_and_nop();
            } else {
                // Enable interrupts, then halt CPU (to save power) until the next interrupt is actually fired.
                interrupt::enable_and_halt();
            }
        }
    }
}

#[no_mangle]
pub extern fn kmain_ap(id: usize) {
    CPU_ID.store(id, Ordering::SeqCst);

    context::init();

    let pid = syscall::getpid();
    println!("AP {}: {:?}", id, pid);

    loop {
        unsafe { interrupt::enable_and_halt() }
    }
}
<|endoftext|>


<commit_before>import std::time;
import std::time::tm;
import m2_request = mongrel2::request;

import router::router;
import request::request;
import response::response;
import middleware::middleware;

import to_bytes::to_bytes;

type mre<T> = {
    m2: mongrel2::connection,
    router: router::router<T>,
    middleware: middleware::middleware<T>,
    mk_data: fn@() -> T,
};

fn mre<T: copy>(m2: mongrel2::connection,
          middleware: middleware::middleware<T>,
          mk_data: fn@() -> T) -> mre<T> {
    {
        m2: m2,
        router: router::router(),
        middleware: middleware,
        mk_data: mk_data,
    }
}

impl mre<T: copy> for mre<T> {
    fn run() {
        io::println(#fmt("Starting up %? -> %?",
            self.m2.sub_addrs(),
            self.m2.pub_addrs()));

        loop {
            let m2_req = self.m2.recv();

            // Ignore close requests for now.
            if m2_req.is_disconnect() { cont; }

            let rep = response::response(self.m2, m2_req);

            let req = alt request::request(m2_req, rep, self.mk_data()) {
              none {
                // Ignore this request if it's malformed.
                cont;
              }
              some(req) { req }
            };

            self.middleware.wrap(req, rep);

            alt self.router.find(req.method, req.path()) {
              none { rep.reply_http(404u, "") }
              some((handler, m)) { handler(req, rep, m) }
            };
        }
    }
}
<commit_msg>These functions were renamed.<commit_after>import std::time;
import std::time::tm;
import m2_request = mongrel2::request;

import router::router;
import request::request;
import response::response;
import middleware::middleware;

import to_bytes::to_bytes;

type mre<T> = {
    m2: mongrel2::connection,
    router: router::router<T>,
    middleware: middleware::middleware<T>,
    mk_data: fn@() -> T,
};

fn mre<T: copy>(m2: mongrel2::connection,
          middleware: middleware::middleware<T>,
          mk_data: fn@() -> T) -> mre<T> {
    {
        m2: m2,
        router: router::router(),
        middleware: middleware,
        mk_data: mk_data,
    }
}

impl mre<T: copy> for mre<T> {
    fn run() {
        io::println(#fmt("Starting up %? -> %?",
            self.m2.req_addrs(),
            self.m2.rep_addrs()));

        loop {
            let m2_req = self.m2.recv();

            // Ignore close requests for now.
            if m2_req.is_disconnect() { cont; }

            let rep = response::response(self.m2, m2_req);

            let req = alt request::request(m2_req, rep, self.mk_data()) {
              none {
                // Ignore this request if it's malformed.
                cont;
              }
              some(req) { req }
            };

            self.middleware.wrap(req, rep);

            alt self.router.find(req.method, req.path()) {
              none { rep.reply_http(404u, "") }
              some((handler, m)) { handler(req, rep, m) }
            };
        }
    }
}
<|endoftext|>