<commit_before>extern crate ws;

use message::PeerMessage;

use std::cmp;
use std::collections::hash_map::HashMap;
use std::sync::{Arc, Mutex};

use bincode::rustc_serialize::decode;

/// Connection handler for incoming messages.
pub struct MessageHandler {
    /// Sender that is used to communicate for the handler.
    pub ws: ws::Sender,
    /// An arc clone of my local vector clocks.
    pub clocks: Arc<Mutex<HashMap<String, u32>>>,
    /// My address/name.
    pub me: String,
}

//TODO: When connected to another client, add it to the `clocks`.
impl ws::Handler for MessageHandler {
    fn on_message(&mut self, msg: ws::Message) -> ws::Result<()> {
        match msg {
            ws::Message::Binary(vector) => {
                let encoded_msg = &*vector.into_boxed_slice();
                let message: PeerMessage = decode(encoded_msg).unwrap();

                // Checkings if all clients have everyone's clocks
                //TODO: Maybe remove this since we add vclocks on connect? Test first!
//                let mut clocks = self.clocks.lock().unwrap();
//                for (key, val) in message.clocks.iter() {
//                    let value = val.clone();
//                    let keyer = key.clone();
//                    if !clocks.contains_key(&keyer) {
//                        clocks.insert(String::from(keyer), value);
//                    }
//                }

                info!("Peer {} with clocks: {:?} got message: {}",
                        message.sender, message.clocks, message.message);
                //message_checking(&mut clocks, message.clone());
				message_handler(message, self.clocks.clone());
                Ok(())
            },
            ws::Message::Text(string) => {
                info!("Received peer's name. Adding {} to the client list", string);
                let mut clocks = self.clocks.lock().unwrap();
                clocks.insert(string, 0u32);
                Ok(())
            }
        }
    }
    fn on_close(&mut self, code: ws::CloseCode, reason: &str) {
        //TODO: Remove disconnected client's clock from our vclock copy.
        if reason.is_empty() {
            info!("Client disconnected with code: {:?}", code); //This works: CloseCode::Abnormal
        } else {
            info!("Client disconnected with code: {:?} and reason: {}", code, reason);
        }
    }
    fn on_error(&mut self, err: ws::Error) {
        warn!("Error family robinson! {:?}", err.kind);
    }
}

pub struct MessageFactory {
    vclocks: Arc<Mutex<HashMap<String, u32>>>,
    me: String,
}

impl MessageFactory {
    pub fn build(vclocks: Arc<Mutex<HashMap<String, u32>>>) -> MessageFactory {
        MessageFactory {
            vclocks: vclocks,
            me: String::from("undefined"),
        }
    }
    pub fn me(&self, me: &str) -> MessageFactory {
        MessageFactory {
            vclocks: self.vclocks.clone(),
            me: String::from(me),
        }
    }
}

impl ws::Factory for MessageFactory {
    type Handler = MessageHandler;

    fn connection_made(&mut self, ws: ws::Sender) -> Self::Handler {
        MessageHandler {
            ws: ws,
            clocks: self.vclocks.clone(),
            me: self.me.clone(),
        }
    }

    fn server_connected(&mut self, ws: ws::Sender) -> Self::Handler {
        ws.send(self.me.clone()).unwrap();
        MessageHandler {
            ws: ws,
            clocks: self.vclocks.clone(),
            me: self.me.clone(),
        }
    }

    fn client_connected(&mut self, ws: ws::Sender) -> Self::Handler {
        ws.send(self.me.clone()).unwrap();
        MessageHandler {
            ws: ws,
            clocks: self.vclocks.clone(),
            me: self.me.clone(),
        }
    }
}

fn message_handler(message: PeerMessage, vclocks: Arc<Mutex<HashMap<String, u32>>>) {
    let mut vclocks = vclocks.lock().unwrap();
    for (key, val) in message.clocks.iter() {
        if let Some(lval) = vclocks.get(&key.clone()) {
            info!("key: {} incoming val: {} and lval: {}", key, val, lval);
            if lval <= val {
                info!("val <= lval");
                continue;
            } else { // There's a clock that is greater than what we have
                // Push to local buffer
                // return
                info!("clock discrepency");
                ()
            }
        }
    }

    // Update clocks
    for (key, val) in message.clocks.iter() {
        let lval = vclocks.get(&key.clone()).unwrap().clone();
        let max_val = cmp::max(val, &lval);
        vclocks.insert(key.clone(), *max_val);
    }
//	macro_rules! hashmap {
//		($( $key: expr => $val: expr ),*) => {{
//			 let mut map = ::std::collections::HashMap::new();
//			 $( map.insert($key, $val); )*
//			 map
//		}}
//	}
    //let map = hashmap!['a' => 0, 'b' => 1];
}

<commit_msg>Add fake delay<commit_after>extern crate ws;

use message::PeerMessage;

use std::cmp;
use std::collections::hash_map::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use std::thread;

use bincode::rustc_serialize::decode;

/// Connection handler for incoming messages.
pub struct MessageHandler {
    /// Sender that is used to communicate for the handler.
    pub ws: ws::Sender,
    /// An arc clone of my local vector clocks.
    pub clocks: Arc<Mutex<HashMap<String, u32>>>,
    /// My address/name.
    pub me: String,
}

//TODO: When connected to another client, add it to the `clocks`.
impl ws::Handler for MessageHandler {
    fn on_message(&mut self, msg: ws::Message) -> ws::Result<()> {
        match msg {
            ws::Message::Binary(vector) => {
                let encoded_msg = &*vector.into_boxed_slice();
                let message: PeerMessage = decode(encoded_msg).unwrap();

                info!("Peer {} with clocks: {:?} got message: {}",
                        message.sender, message.clocks, message.message);
                if message.sender.as_str() == "127.0.0.1:3013" {
                    info!("Faking delay!");
                    thread::sleep(Duration::from_millis(4000))
                }
                //message_checking(&mut clocks, message.clone());
				message_handler(message, self.clocks.clone());
                Ok(())
            },
            ws::Message::Text(string) => {
                info!("Received peer's name. Adding {} to the client list", string);
                let mut clocks = self.clocks.lock().unwrap();
                clocks.insert(string, 0u32);
                Ok(())
            }
        }
    }
    fn on_close(&mut self, code: ws::CloseCode, reason: &str) {
        //TODO: Remove disconnected client's clock from our vclock copy.
        if reason.is_empty() {
            info!("Client disconnected with code: {:?}", code); //This works: CloseCode::Abnormal
        } else {
            info!("Client disconnected with code: {:?} and reason: {}", code, reason);
        }
    }
    fn on_error(&mut self, err: ws::Error) {
        warn!("Error family robinson! {:?}", err.kind);
    }
}

pub struct MessageFactory {
    vclocks: Arc<Mutex<HashMap<String, u32>>>,
    me: String,
}

impl MessageFactory {
    pub fn build(vclocks: Arc<Mutex<HashMap<String, u32>>>) -> MessageFactory {
        MessageFactory {
            vclocks: vclocks,
            me: String::from("undefined"),
        }
    }
    pub fn me(&self, me: &str) -> MessageFactory {
        MessageFactory {
            vclocks: self.vclocks.clone(),
            me: String::from(me),
        }
    }
}

impl ws::Factory for MessageFactory {
    type Handler = MessageHandler;

    fn connection_made(&mut self, ws: ws::Sender) -> Self::Handler {
        MessageHandler {
            ws: ws,
            clocks: self.vclocks.clone(),
            me: self.me.clone(),
        }
    }

    fn server_connected(&mut self, ws: ws::Sender) -> Self::Handler {
        ws.send(self.me.clone()).unwrap();
        MessageHandler {
            ws: ws,
            clocks: self.vclocks.clone(),
            me: self.me.clone(),
        }
    }

    fn client_connected(&mut self, ws: ws::Sender) -> Self::Handler {
        ws.send(self.me.clone()).unwrap();
        MessageHandler {
            ws: ws,
            clocks: self.vclocks.clone(),
            me: self.me.clone(),
        }
    }
}

fn message_handler(message: PeerMessage, vclocks: Arc<Mutex<HashMap<String, u32>>>) {
    let mut vclocks = vclocks.lock().unwrap();
    for (key, val) in message.clocks.iter() {
        if let Some(lval) = vclocks.get(&key.clone()) {
            info!("key: {} incoming val: {} and lval: {}", key, val, lval);
            if lval <= val {
                info!("val <= lval");
                continue;
            } else { // There's a clock that is greater than what we have
                // Push to local buffer
                // return
                info!("clock discrepency");
                ()
            }
        }
    }

    // Update clocks
    for (key, val) in message.clocks.iter() {
        let lval = vclocks.get(&key.clone()).unwrap().clone();
        let max_val = cmp::max(val, &lval);
        vclocks.insert(key.clone(), *max_val);
    }
//	macro_rules! hashmap {
//		($( $key: expr => $val: expr ),*) => {{
//			 let mut map = ::std::collections::HashMap::new();
//			 $( map.insert($key, $val); )*
//			 map
//		}}
//	}
    //let map = hashmap!['a' => 0, 'b' => 1];
}

<|endoftext|>


<commit_before>#[macro_use]
extern crate log;
extern crate log4rs;
extern crate iron;
extern crate logger;
extern crate router;
extern crate persistent;
extern crate r2d2;
extern crate r2d2_postgres;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate potboiler_common;
extern crate serde_json;
extern crate hyper;
extern crate hybrid_clocks;
mod tables;

use iron::prelude::*;
use iron::status;
use logger::Logger;
use persistent::Read as PRead;
use persistent::State;
use potboiler_common::{db, iron_str_error, server_id};
use potboiler_common::string_error::StringError;
use potboiler_common::types::{CRDT, LWW, Log};
use r2d2_postgres::PostgresConnectionManager;
use router::Router;
use std::env;
use std::io::Read;
use std::ops::Deref;

pub type PostgresConnection = r2d2::PooledConnection<PostgresConnectionManager>;

lazy_static! {
    static ref SERVER_URL: String = env::var("SERVER_URL").expect("Needed SERVER_URL");
}

include!(concat!(env!("OUT_DIR"), "/serde_types.rs"));

fn get_key(_: &mut Request) -> IronResult<Response> {
    Ok(Response::with((status::Ok, "get_key")))
}

fn update_key(req: &mut Request) -> IronResult<Response> {
    let body_string = {
        let mut body = String::new();
        req.body.read_to_string(&mut body).expect("could read from body");
        body
    };
    let mut json: serde_json::Value = match serde_json::de::from_str(&body_string) {
        Ok(val) => val,
        Err(err) => return Err(IronError::new(err, (status::BadRequest, "Bad JSON"))),
    };
    let map = json.as_object_mut().unwrap();
    map.insert("table".to_string(),
               serde_json::to_value(&potboiler_common::get_req_key(req, "table").unwrap()));
    map.insert("key".to_string(),
               serde_json::to_value(&potboiler_common::get_req_key(req, "key").unwrap()));

    let client = hyper::client::Client::new();
    let res = client.post(SERVER_URL.deref())
        .body(&serde_json::ser::to_string(map).unwrap())
        .send()
        .expect("sender ok");
    assert_eq!(res.status, hyper::status::StatusCode::Created);
    Ok(Response::with((status::Ok, "update_key")))
}

fn make_table(conn: &PostgresConnection, table_name: &str, kind: &CRDT) -> IronResult<()> {
    match kind {
        &CRDT::LWW => {
            conn.execute(&format!("CREATE TABLE IF NOT EXISTS {} (key VARCHAR(1024) PRIMARY KEY, value \
                                   JSONB, crdt JSONB)",
                                  &table_name),
                         &[])
                .map_err(iron_str_error)?;
        }
        &CRDT::ORSET => {
            conn.execute(&format!("CREATE TABLE IF NOT EXISTS {} (key VARCHAR(1024) PRIMARY KEY, crdt \
                                   JSONB)",
                                  &table_name),
                         &[])
                .map_err(iron_str_error)?;
            conn.execute(&format!("CREATE TABLE IF NOT EXISTS {}_items (\
                                   key VARCHAR(1024), \
                                   item VARCHAR(1024), PRIMARY KEY(key, item))",
                                  &table_name),
                         &[])
                .map_err(iron_str_error)?;
        }
        &CRDT::GSET => {
            error!("No G-Set make table yet");
        }
    }
    return Ok(());
}

fn string_iron_error(error: &str) -> IronResult<Response> {
    Err(IronError::new(StringError(error.to_string()), (status::BadRequest, error)))
}

fn get_crdt(conn: &PostgresConnection,
            table: &String,
            key: &String)
            -> IronResult<Option<serde_json::Value>> {
    let stmt = conn.prepare(&format!("select crdt from {} where key=$1", table)).map_err(iron_str_error)?;
    let results = stmt.query(&[&key]).map_err(iron_str_error)?;
    if results.is_empty() {
        Ok(None)
    } else if results.len() == 1 {
        let row = results.get(0);
        let raw_crdt: serde_json::Value = row.get("crdt");
        Ok(Some(raw_crdt))
    } else {
        Err(iron_str_error(StringError::from(format!("{} entries for key {}", results.len(), key))))
    }
}

fn new_event(req: &mut Request) -> IronResult<Response> {
    let body_string = {
        let mut body = String::new();
        req.body.read_to_string(&mut body).expect("could read from body");
        body
    };
    let json: serde_json::Value = match serde_json::de::from_str(&body_string) {
        Ok(val) => val,
        Err(err) => return Err(IronError::new(err, (status::BadRequest, "Bad JSON"))),
    };
    info!("body: {:?}", json);
    let log = try!(serde_json::from_value::<Log>(json).map_err(iron_str_error));
    info!("log: {:?}", log);
    let change: Change = serde_json::from_value(log.data).unwrap();
    info!("change: {:?}", change);
    let tables = tables::get_tables(req);
    match tables.get(&change.table) {
        None => return string_iron_error("Can't find table"),
        Some(&val) => {
            match val {
                CRDT::LWW => {
                    match change.op {
                        Operation::Set => {
                            let crdt_to_use: Option<CRDT> = if &change.table == tables::CONFIG_TABLE {
                                let crdt_type: &str = change.change
                                    .find("crdt")
                                    .ok_or(StringError::from("No CRDT key"))
                                    .map_err(iron_str_error)?
                                    .as_str()
                                    .ok_or(StringError::from("CRDT value isn't string!"))
                                    .map_err(iron_str_error)?;
                                // FIXME: format! bit is a hacky workaround for
                                // https://github.com/serde-rs/serde/issues/251
                                serde_json::from_str(&format!("\"{}\"", crdt_type)).map_err(iron_str_error)?
                            } else {
                                None
                            };
                            let conn = get_pg_connection!(&req);
                            let raw_crdt = get_crdt(&conn, &change.table, &change.key)?;
                            match raw_crdt {
                                None => {
                                    let lww = LWW {
                                        when: log.when,
                                        data: change.change.clone(),
                                    };
                                    conn.execute(&format!("INSERT INTO {} (key, value, crdt) VALUES ($1, \
                                                           $2, $3)",
                                                          &change.table),
                                                 &[&change.key, &change.change, &serde_json::to_value(&lww)])
                                        .expect("insert worked");
                                    if &change.table == tables::CONFIG_TABLE {
                                        let crdt = crdt_to_use.unwrap();
                                        make_table(&conn, &change.key, &crdt)?;
                                        tables::add_table(req, &change.key, &crdt);
                                    }
                                }
                                Some(raw_crdt) => {
                                    let mut lww: LWW = serde_json::from_value(raw_crdt)
                                        .expect("bad raw crdt");
                                    if lww.when < log.when {
                                        lww.when = log.when;
                                        lww.data = change.change.clone();
                                        conn.execute(&format!("UPDATE {} set value=$2, crdt=$3 where \
                                                               key=$1",
                                                              &change.table),
                                                     &[&change.key,
                                                       &change.change,
                                                       &serde_json::to_value(&lww)])
                                            .expect("update worked");
                                    } else {
                                        info!("Earlier event, skipping");
                                    }
                                }
                            }
                        }
                        _ => {
                            return string_iron_error("LWW only supports Set, not Add/Remove");
                        }
                    }
                }
                CRDT::ORSET => {
                    let op: ORSetOp = serde_json::from_value(change.change).map_err(iron_str_error)?;
                    let conn = get_pg_connection!(&req);
                    let raw_crdt = get_crdt(&conn, &change.table, &change.key)?;
                    let (mut crdt, existing) = match raw_crdt {
                        Some(val) => (serde_json::from_value(val).map_err(iron_str_error)?, true),
                        None => {
                            (ORSet {
                                 adds: HashMap::new(),
                                 removes: HashMap::new(),
                             },
                             false)
                        }
                    };
                    let trans = conn.transaction().unwrap();
                    match change.op {
                        Operation::Add => {
                            if !crdt.removes.contains_key(&op.key) && !crdt.adds.contains_key(&op.key) {
                                trans.execute(&format!("INSERT INTO {}_items (key, item) VALUES ($1, $2)",
                                                      &change.table),
                                             &[&op.key, &op.item])
                                    .map_err(iron_str_error)?;
                                crdt.adds.insert(op.key, op.item);
                            }
                        }
                        Operation::Remove => {
                            trans.execute(&format!("DELETE FROM {}_items where key=$1", &change.table),
                                         &[&op.key])
                                .map_err(iron_str_error)?;
                            crdt.adds.remove(&op.key);
                            crdt.removes.insert(op.key, op.item);
                        }
                        _ => {
                            return string_iron_error("ORSET only supports Add/Remove");
                        }
                    }
                    debug!("OR-Set for {}: {:?}", &change.table, &crdt);
                    if existing {
                        trans.execute(&format!("UPDATE {} set crdt=$2 where key=$1", &change.table),
                                     &[&change.key, &serde_json::to_value(&crdt)])
                            .map_err(iron_str_error)?;
                    } else {
                        trans.execute(&format!("INSERT INTO {} (key, crdt) VALUES ($1, $2)", &change.table),
                                     &[&change.key, &serde_json::to_value(&crdt)])
                            .map_err(iron_str_error)?;
                    }
                    let mut items: Vec<&str> = Vec::new();
                    for key in crdt.adds.keys() {
                        if crdt.removes.contains_key(key) {
                            continue;
                        }
                        items.push(crdt.adds.get(key).unwrap());
                    }
                    debug!("Items: {:?}", items);
                    trans.commit().map_err(iron_str_error)?;
                }
                _ => {
                    return string_iron_error("Only support LWW and OR-Set so far");
                }
            }
        }
    }
    Ok(Response::with(status::NoContent))
}

fn main() {
    log4rs::init_file("log.yaml", Default::default()).expect("log config ok");
    let client = hyper::client::Client::new();

    let mut map: serde_json::Map<String, String> = serde_json::Map::new();
    let host: &str = &env::var("HOST").unwrap_or("localhost".to_string());
    map.insert("url".to_string(),
               format!("http://{}:8001/kv/event", host).to_string());
    let res = client.post(&format!("{}/register", SERVER_URL.deref()))
        .body(&serde_json::ser::to_string(&map).unwrap())
        .send()
        .expect("Register ok");
    assert_eq!(res.status, hyper::status::StatusCode::NoContent);

    let db_url: &str = &env::var("DATABASE_URL").expect("Needed DATABASE_URL");
    let pool = db::get_pool(db_url);
    let conn = pool.get().unwrap();
    match make_table(&conn, "_config", &CRDT::LWW) {
        Ok(_) => {}
        Err(err) => {
            error!("Error while making config table: {}", err);
            return;
        }
    }
    let (logger_before, logger_after) = Logger::new(None);
    let mut router = Router::new();
    router.get("/kv/:table/:key", get_key);
    router.post("/kv/:table/:key", update_key);
    router.post("/kv/event", new_event);
    let mut chain = Chain::new(router);
    chain.link_before(logger_before);
    chain.link_after(logger_after);
    chain.link_before(PRead::<server_id::ServerId>::one(server_id::setup()));
    chain.link(PRead::<db::PostgresDB>::both(pool));
    let tables = tables::init_tables(&conn);
    chain.link(State::<tables::Tables>::both(tables));
    info!("Potboiler-kv booted");
    Iron::new(chain).http("0.0.0.0:8001").unwrap();
}
<commit_msg>Reduce kv indent<commit_after>#[macro_use]
extern crate log;
extern crate log4rs;
extern crate iron;
extern crate logger;
extern crate router;
extern crate persistent;
extern crate r2d2;
extern crate r2d2_postgres;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate potboiler_common;
extern crate serde_json;
extern crate hyper;
extern crate hybrid_clocks;
mod tables;

use iron::prelude::*;
use iron::status;
use logger::Logger;
use persistent::Read as PRead;
use persistent::State;
use potboiler_common::{db, iron_str_error, server_id};
use potboiler_common::string_error::StringError;
use potboiler_common::types::{CRDT, LWW, Log};
use r2d2_postgres::PostgresConnectionManager;
use router::Router;
use std::env;
use std::io::Read;
use std::ops::Deref;

pub type PostgresConnection = r2d2::PooledConnection<PostgresConnectionManager>;

lazy_static! {
    static ref SERVER_URL: String = env::var("SERVER_URL").expect("Needed SERVER_URL");
}

include!(concat!(env!("OUT_DIR"), "/serde_types.rs"));

fn get_key(_: &mut Request) -> IronResult<Response> {
    Ok(Response::with((status::Ok, "get_key")))
}

fn update_key(req: &mut Request) -> IronResult<Response> {
    let body_string = {
        let mut body = String::new();
        req.body.read_to_string(&mut body).expect("could read from body");
        body
    };
    let mut json: serde_json::Value = match serde_json::de::from_str(&body_string) {
        Ok(val) => val,
        Err(err) => return Err(IronError::new(err, (status::BadRequest, "Bad JSON"))),
    };
    let map = json.as_object_mut().unwrap();
    map.insert("table".to_string(),
               serde_json::to_value(&potboiler_common::get_req_key(req, "table").unwrap()));
    map.insert("key".to_string(),
               serde_json::to_value(&potboiler_common::get_req_key(req, "key").unwrap()));

    let client = hyper::client::Client::new();
    let res = client.post(SERVER_URL.deref())
        .body(&serde_json::ser::to_string(map).unwrap())
        .send()
        .expect("sender ok");
    assert_eq!(res.status, hyper::status::StatusCode::Created);
    Ok(Response::with((status::Ok, "update_key")))
}

fn make_table(conn: &PostgresConnection, table_name: &str, kind: &CRDT) -> IronResult<()> {
    match kind {
        &CRDT::LWW => {
            conn.execute(&format!("CREATE TABLE IF NOT EXISTS {} (key VARCHAR(1024) PRIMARY KEY, value \
                                   JSONB, crdt JSONB)",
                                  &table_name),
                         &[])
                .map_err(iron_str_error)?;
        }
        &CRDT::ORSET => {
            conn.execute(&format!("CREATE TABLE IF NOT EXISTS {} (key VARCHAR(1024) PRIMARY KEY, crdt \
                                   JSONB)",
                                  &table_name),
                         &[])
                .map_err(iron_str_error)?;
            conn.execute(&format!("CREATE TABLE IF NOT EXISTS {}_items (\
                                   key VARCHAR(1024), \
                                   item VARCHAR(1024), PRIMARY KEY(key, item))",
                                  &table_name),
                         &[])
                .map_err(iron_str_error)?;
        }
        &CRDT::GSET => {
            error!("No G-Set make table yet");
        }
    }
    return Ok(());
}

fn string_iron_error(error: &str) -> IronResult<Response> {
    Err(IronError::new(StringError(error.to_string()), (status::BadRequest, error)))
}

fn get_crdt(conn: &PostgresConnection,
            table: &String,
            key: &String)
            -> IronResult<Option<serde_json::Value>> {
    let stmt = conn.prepare(&format!("select crdt from {} where key=$1", table)).map_err(iron_str_error)?;
    let results = stmt.query(&[&key]).map_err(iron_str_error)?;
    if results.is_empty() {
        Ok(None)
    } else if results.len() == 1 {
        let row = results.get(0);
        let raw_crdt: serde_json::Value = row.get("crdt");
        Ok(Some(raw_crdt))
    } else {
        Err(iron_str_error(StringError::from(format!("{} entries for key {}", results.len(), key))))
    }
}

fn new_event(req: &mut Request) -> IronResult<Response> {
    let body_string = {
        let mut body = String::new();
        req.body.read_to_string(&mut body).expect("could read from body");
        body
    };
    let json: serde_json::Value = match serde_json::de::from_str(&body_string) {
        Ok(val) => val,
        Err(err) => return Err(IronError::new(err, (status::BadRequest, "Bad JSON"))),
    };
    info!("body: {:?}", json);
    let log = try!(serde_json::from_value::<Log>(json).map_err(iron_str_error));
    info!("log: {:?}", log);
    let change: Change = serde_json::from_value(log.data).unwrap();
    info!("change: {:?}", change);
    let tables = tables::get_tables(req);
    let table_type = match tables.get(&change.table) {
        None => return string_iron_error("Can't find table"),
        Some(&val) => val,
    };
    match table_type {
        CRDT::LWW => {
            match change.op {
                Operation::Set => {
                    let crdt_to_use: Option<CRDT> = if &change.table == tables::CONFIG_TABLE {
                        let crdt_type: &str = change.change
                            .find("crdt")
                            .ok_or(StringError::from("No CRDT key"))
                            .map_err(iron_str_error)?
                            .as_str()
                            .ok_or(StringError::from("CRDT value isn't string!"))
                            .map_err(iron_str_error)?;
                        // FIXME: format! bit is a hacky workaround for
                        // https://github.com/serde-rs/serde/issues/251
                        serde_json::from_str(&format!("\"{}\"", crdt_type)).map_err(iron_str_error)?
                    } else {
                        None
                    };
                    let conn = get_pg_connection!(&req);
                    let raw_crdt = get_crdt(&conn, &change.table, &change.key)?;
                    match raw_crdt {
                        None => {
                            let lww = LWW {
                                when: log.when,
                                data: change.change.clone(),
                            };
                            conn.execute(&format!("INSERT INTO {} (key, value, crdt) VALUES ($1, $2, $3)",
                                                  &change.table),
                                         &[&change.key, &change.change, &serde_json::to_value(&lww)])
                                .expect("insert worked");
                            if &change.table == tables::CONFIG_TABLE {
                                let crdt = crdt_to_use.unwrap();
                                make_table(&conn, &change.key, &crdt)?;
                                tables::add_table(req, &change.key, &crdt);
                            }
                        }
                        Some(raw_crdt) => {
                            let mut lww: LWW = serde_json::from_value(raw_crdt).expect("bad raw crdt");
                            if lww.when < log.when {
                                lww.when = log.when;
                                lww.data = change.change.clone();
                                conn.execute(&format!("UPDATE {} set value=$2, crdt=$3 where key=$1",
                                                      &change.table),
                                             &[&change.key, &change.change, &serde_json::to_value(&lww)])
                                    .expect("update worked");
                            } else {
                                info!("Earlier event, skipping");
                            }
                        }
                    }
                }
                _ => {
                    return string_iron_error("LWW only supports Set, not Add/Remove");
                }
            }
        }
        CRDT::ORSET => {
            let op: ORSetOp = serde_json::from_value(change.change).map_err(iron_str_error)?;
            let conn = get_pg_connection!(&req);
            let raw_crdt = get_crdt(&conn, &change.table, &change.key)?;
            let (mut crdt, existing) = match raw_crdt {
                Some(val) => (serde_json::from_value(val).map_err(iron_str_error)?, true),
                None => {
                    (ORSet {
                         adds: HashMap::new(),
                         removes: HashMap::new(),
                     },
                     false)
                }
            };
            let trans = conn.transaction().unwrap();
            match change.op {
                Operation::Add => {
                    if !crdt.removes.contains_key(&op.key) && !crdt.adds.contains_key(&op.key) {
                        trans.execute(&format!("INSERT INTO {}_items (key, item) VALUES ($1, $2)",
                                              &change.table),
                                     &[&op.key, &op.item])
                            .map_err(iron_str_error)?;
                        crdt.adds.insert(op.key, op.item);
                    }
                }
                Operation::Remove => {
                    trans.execute(&format!("DELETE FROM {}_items where key=$1", &change.table),
                                 &[&op.key])
                        .map_err(iron_str_error)?;
                    crdt.adds.remove(&op.key);
                    crdt.removes.insert(op.key, op.item);
                }
                _ => {
                    return string_iron_error("ORSET only supports Add/Remove");
                }
            }
            debug!("OR-Set for {}: {:?}", &change.table, &crdt);
            if existing {
                trans.execute(&format!("UPDATE {} set crdt=$2 where key=$1", &change.table),
                             &[&change.key, &serde_json::to_value(&crdt)])
                    .map_err(iron_str_error)?;
            } else {
                trans.execute(&format!("INSERT INTO {} (key, crdt) VALUES ($1, $2)", &change.table),
                             &[&change.key, &serde_json::to_value(&crdt)])
                    .map_err(iron_str_error)?;
            }
            let mut items: Vec<&str> = Vec::new();
            for key in crdt.adds.keys() {
                if crdt.removes.contains_key(key) {
                    continue;
                }
                items.push(crdt.adds.get(key).unwrap());
            }
            debug!("Items: {:?}", items);
            trans.commit().map_err(iron_str_error)?;
        }
        _ => {
            return string_iron_error("Only support LWW and OR-Set so far");
        }
    }
    Ok(Response::with(status::NoContent))
}

fn main() {
    log4rs::init_file("log.yaml", Default::default()).expect("log config ok");
    let client = hyper::client::Client::new();

    let mut map: serde_json::Map<String, String> = serde_json::Map::new();
    let host: &str = &env::var("HOST").unwrap_or("localhost".to_string());
    map.insert("url".to_string(),
               format!("http://{}:8001/kv/event", host).to_string());
    let res = client.post(&format!("{}/register", SERVER_URL.deref()))
        .body(&serde_json::ser::to_string(&map).unwrap())
        .send()
        .expect("Register ok");
    assert_eq!(res.status, hyper::status::StatusCode::NoContent);

    let db_url: &str = &env::var("DATABASE_URL").expect("Needed DATABASE_URL");
    let pool = db::get_pool(db_url);
    let conn = pool.get().unwrap();
    match make_table(&conn, "_config", &CRDT::LWW) {
        Ok(_) => {}
        Err(err) => {
            error!("Error while making config table: {}", err);
            return;
        }
    }
    let (logger_before, logger_after) = Logger::new(None);
    let mut router = Router::new();
    router.get("/kv/:table/:key", get_key);
    router.post("/kv/:table/:key", update_key);
    router.post("/kv/event", new_event);
    let mut chain = Chain::new(router);
    chain.link_before(logger_before);
    chain.link_after(logger_after);
    chain.link_before(PRead::<server_id::ServerId>::one(server_id::setup()));
    chain.link(PRead::<db::PostgresDB>::both(pool));
    let tables = tables::init_tables(&conn);
    chain.link(State::<tables::Tables>::both(tables));
    info!("Potboiler-kv booted");
    Iron::new(chain).http("0.0.0.0:8001").unwrap();
}
<|endoftext|>


<commit_before>extern crate log;
extern crate discotech_zookeeper;

use config::*;

use rustc_serialize::json;
use std::sync::RwLock;
use std::time::Duration;
use std::collections::HashMap;
use std::collections::HashSet;
use discotech_zookeeper::{Acl, CreateMode, Watcher, WatchedEvent, ZkError, ZooKeeper};
use discotech_zookeeper::perms;

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct ServiceEndpoint {
  pub host: String,
  pub port: u16,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct ServersetMember {
  pub serviceEndpoint: ServiceEndpoint,
  pub additionalEndpoints: HashMap<String, ServiceEndpoint>,
  pub status: String,
}

struct NullWatcher;
impl Watcher for NullWatcher {
  fn handle(&self, e: &WatchedEvent) {}
}

pub struct Serverset {
  config: DiscoConfig,
  zk_client: ZooKeeper,
  pub members: RwLock<HashMap<String, ServersetMember>>,
}
impl Serverset {
  pub fn new(discoConfig: DiscoConfig) -> Serverset {
    match ZooKeeper::connect(format!("{}:{}/",
        discoConfig.zookeeper_host, discoConfig.zookeeper_port).as_str(),
        Duration::from_secs(discoConfig.zookeeper_timeout_secs), NullWatcher) {
      Err(reason) => panic!("Unable to connect to ZooKeeper: {}", reason),
      Ok(client) => Serverset{
        config: discoConfig,
        zk_client: client,
        members: RwLock::new(HashMap::new()),
      },
    }
  }

  fn remove_member(&self, member_znode: &String) {
    self.members.write().unwrap().remove(member_znode);
  }

  fn update_member(&self, member_znode: &String) {
    debug!("Adding Serverset member: {}", member_znode);

    // Reads Serverset member's ZNode data and attempts to parse it into a String.
    let member_json_opt = match self.zk_client.get_data(format!("{}/{}",
        self.config.serverset_znode, member_znode).as_str(), false) {
      Err(reason) => {
        error!("Could not obtain node data for {} from ZooKeeper: {}", member_znode,
            reason);
        None
      },
      Ok(node_data) => match String::from_utf8(node_data.0) {
        Err(reason) => {
          error!("Could not parse node string: {}", reason);
          None
        },
        Ok(node_string) => Some(node_string),
      },
    };

    // Attempts to parse Serverset member's ZNode into a ServersetMember struct.
    let member_opt: Option<ServersetMember> = match member_json_opt {
      None => None,
      Some(member_json) => match json::decode(&member_json) {
        Err(reason) => {
          error!("Could not parse node JSON: {}", reason);
          None
        },
        Ok(member) => Some(member),
      },
    };

    // If all has gone well, grabs a write lock on the members HashMap and updates it with
    // the newly-unwrapped ServersetMember.
    match member_opt {
      None => None,
      Some(member) => match member.status.as_ref() {
        "ALIVE" => {
          self.members.write().unwrap().insert(member_znode.clone(), member)
        },
        _ => None,
      },
    };
  }

  pub fn update_members(&self) {
    // Reconciles our local representation of the Serverset with that which has been
    // stored in ZooKeeper.
    debug!("Updating Serverset members...");
    match self.zk_client.get_children(self.config.serverset_znode.as_str(), false) {
      Err(reason) => error!("Unable to get children for {}: {}",
          self.config.serverset_znode, reason),
      Ok(serverset_children) => {
        debug!("Children: {:?}", serverset_children);
        // Updates all serverset members in parallel, tracking those which we've seen.
        let mut current_member_znodes = HashSet::new();
        for current_member_znode in serverset_children.iter() {
          current_member_znodes.insert(current_member_znode);
          self.update_member(current_member_znode);
        }
        // Removes all members that have dropped out of the serverset.
        for old_member_znode in self.members.read().unwrap().keys() {
          if !current_member_znodes.contains(old_member_znode) {
            self.remove_member(old_member_znode);
          }
        }
      },
    }
  }
}
<commit_msg>No more thread stack overflows<commit_after>extern crate log;
extern crate discotech_zookeeper;

use config::*;

use rustc_serialize::json;
use std::sync::RwLock;
use std::time::Duration;
use std::collections::HashMap;
use std::collections::HashSet;
use discotech_zookeeper::{Acl, CreateMode, Watcher, WatchedEvent, ZkError, ZooKeeper};
use discotech_zookeeper::perms;

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct ServiceEndpoint {
  pub host: String,
  pub port: u16,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct ServersetMember {
  pub serviceEndpoint: ServiceEndpoint,
  pub additionalEndpoints: HashMap<String, ServiceEndpoint>,
  pub status: String,
}

struct NullWatcher;
impl Watcher for NullWatcher {
  fn handle(&self, e: &WatchedEvent) {}
}

pub struct Serverset {
  config: DiscoConfig,
  zk_client: ZooKeeper,
  pub members: RwLock<HashMap<String, ServersetMember>>,
}
impl Serverset {
  pub fn new(discoConfig: DiscoConfig) -> Serverset {
    match ZooKeeper::connect(format!("{}:{}/",
        discoConfig.zookeeper_host, discoConfig.zookeeper_port).as_str(),
        Duration::from_secs(discoConfig.zookeeper_timeout_secs), NullWatcher) {
      Err(reason) => panic!("Unable to connect to ZooKeeper: {}", reason),
      Ok(client) => Serverset{
        config: discoConfig,
        zk_client: client,
        members: RwLock::new(HashMap::new()),
      },
    }
  }

  fn remove_member(&self, member_znode: &String) {
    self.members.write().unwrap().remove(member_znode);
  }

  fn znode_exists(&self, znode: &String) -> bool {
    match self.zk_client.exists(znode, false) {
      Ok(_) => true,
      _ => false,
    }
  }

  fn update_member(&self, member_znode: &String) {
    debug!("Adding Serverset member: {}", member_znode);
    // If the Serverset member's ZNode does not exist, does not update the member.
    let full_member_znode = format!("{}/{}", self.config.serverset_znode, member_znode);
    if !self.znode_exists(&full_member_znode) {
      return
    }
    // Reads Serverset member's ZNode data and attempts to parse it into a String.
    let member_json_opt = match self.zk_client.get_data(full_member_znode.as_str(),
        false) {
      Err(reason) => {
        error!("Could not obtain node data for {} from ZooKeeper: {}", member_znode,
            reason);
        None
      },
      Ok(node_data) => match String::from_utf8(node_data.0) {
        Err(reason) => {
          error!("Could not parse node string: {}", reason);
          None
        },
        Ok(node_string) => Some(node_string),
      },
    };

    // Attempts to parse Serverset member's ZNode into a ServersetMember struct.
    let member_opt: Option<ServersetMember> = match member_json_opt {
      None => None,
      Some(member_json) => match json::decode(&member_json) {
        Err(reason) => {
          error!("Could not parse node JSON: {}", reason);
          None
        },
        Ok(member) => Some(member),
      },
    };

    // If all has gone well, grabs a write lock on the members HashMap and updates it with
    // the newly-unwrapped ServersetMember.
    match member_opt {
      None => None,
      Some(member) => match member.status.as_ref() {
        "ALIVE" => {
          self.members.write().unwrap().insert(member_znode.clone(), member)
        },
        _ => None,
      },
    };
  }

  pub fn update_members(&self) {
    // Reconciles our local representation of the Serverset with that which has been
    // stored in ZooKeeper.
    debug!("Updating Serverset members...");
    if !self.znode_exists(&self.config.serverset_znode) {
      error!("Could not find Serverset ZNode: {}", self.config.serverset_znode);
      return
    }
    match self.zk_client.get_children(self.config.serverset_znode.as_str(), false) {
      Err(reason) => error!("Unable to get children for {}: {}",
          self.config.serverset_znode, reason),
      Ok(serverset_children) => {
        debug!("Children: {:?}", serverset_children);
        // Updates all serverset members in parallel, tracking those which we've seen.
        let mut current_member_znodes = HashSet::new();
        for current_member_znode in serverset_children.iter() {
          current_member_znodes.insert(current_member_znode);
          self.update_member(current_member_znode);
        }
        // Removes all members that have dropped out of the serverset.
        for old_member_znode in self.members.read().unwrap().keys() {
          if !current_member_znodes.contains(old_member_znode) {
            self.remove_member(old_member_znode);
          }
        }
      },
    }
  }
}
<|endoftext|>


<commit_before>// Copyright 2019 The CryptoCorrosion Contributors
// Copyright 2020 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The ChaCha random number generator.

use ppv_lite86::{dispatch, dispatch_light128};

#[cfg(feature = "rustcrypto_api")]
pub use stream_cipher::generic_array;

pub use ppv_lite86::Machine;
use ppv_lite86::{vec128_storage, ArithOps, BitOps32, LaneWords4, MultiLane, StoreBytes, Vec4};

pub(crate) const BLOCK: usize = 64;
pub(crate) const BLOCK64: u64 = BLOCK as u64;
const LOG2_BUFBLOCKS: u64 = 2;
const BUFBLOCKS: u64 = 1 << LOG2_BUFBLOCKS;
pub(crate) const BUFSZ64: u64 = BLOCK64 * BUFBLOCKS;
pub(crate) const BUFSZ: usize = BUFSZ64 as usize;

#[derive(Clone)]
pub struct ChaCha {
    pub(crate) b: vec128_storage,
    pub(crate) c: vec128_storage,
    pub(crate) d: vec128_storage,
}

#[derive(Clone)]
pub struct State<V> {
    pub(crate) a: V,
    pub(crate) b: V,
    pub(crate) c: V,
    pub(crate) d: V,
}

#[inline(always)]
pub(crate) fn round<V: ArithOps + BitOps32>(mut x: State<V>) -> State<V> {
    x.a += x.b;
    x.d = (x.d ^ x.a).rotate_each_word_right16();
    x.c += x.d;
    x.b = (x.b ^ x.c).rotate_each_word_right20();
    x.a += x.b;
    x.d = (x.d ^ x.a).rotate_each_word_right24();
    x.c += x.d;
    x.b = (x.b ^ x.c).rotate_each_word_right25();
    x
}

#[inline(always)]
pub(crate) fn diagonalize<V: LaneWords4>(mut x: State<V>) -> State<V> {
    x.b = x.b.shuffle_lane_words3012();
    x.c = x.c.shuffle_lane_words2301();
    x.d = x.d.shuffle_lane_words1230();
    x
}
#[inline(always)]
pub(crate) fn undiagonalize<V: LaneWords4>(mut x: State<V>) -> State<V> {
    x.b = x.b.shuffle_lane_words1230();
    x.c = x.c.shuffle_lane_words2301();
    x.d = x.d.shuffle_lane_words3012();
    x
}

impl ChaCha {
    #[inline(always)]
    pub fn new(key: &[u8; 32], nonce: &[u8]) -> Self {
        init_chacha(key, nonce)
    }

    #[inline(always)]
    fn pos64<M: Machine>(&self, m: M) -> u64 {
        let d: M::u32x4 = m.unpack(self.d);
        ((d.extract(1) as u64) << 32) | d.extract(0) as u64
    }

    /// Set 64-bit block count, affecting next refill.
    #[inline(always)]
    pub(crate) fn seek64<M: Machine>(&mut self, m: M, blockct: u64) {
        let d: M::u32x4 = m.unpack(self.d);
        self.d = d
            .insert((blockct >> 32) as u32, 1)
            .insert(blockct as u32, 0)
            .into();
    }

    /// Set 32-bit block count, affecting next refill.
    #[inline(always)]
    pub(crate) fn seek32<M: Machine>(&mut self, m: M, blockct: u32) {
        let d: M::u32x4 = m.unpack(self.d);
        self.d = d.insert(blockct, 0).into();
    }

    /// Produce output from the current state.
    #[inline(always)]
    fn output_narrow<M: Machine>(&mut self, m: M, x: State<M::u32x4>, out: &mut [u8; BLOCK]) {
        let k = m.vec([0x6170_7865, 0x3320_646e, 0x7962_2d32, 0x6b20_6574]);
        (x.a + k).write_le(&mut out[0..16]);
        (x.b + m.unpack(self.b)).write_le(&mut out[16..32]);
        (x.c + m.unpack(self.c)).write_le(&mut out[32..48]);
        (x.d + m.unpack(self.d)).write_le(&mut out[48..64]);
    }

    /// Add one to the block counter (no overflow check).
    #[inline(always)]
    fn inc_block_ct<M: Machine>(&mut self, m: M) {
        let mut pos = self.pos64(m);
        let d0: M::u32x4 = m.unpack(self.d);
        pos += 1;
        let d1 = d0.insert((pos >> 32) as u32, 1).insert(pos as u32, 0);
        self.d = d1.into();
    }

    /// Produce 4 blocks of output, advancing the state
    #[inline(always)]
    pub fn refill4(&mut self, drounds: u32, out: &mut [u8; BUFSZ]) {
        refill_wide(self, drounds, out)
    }

    /// Produce a block of output, advancing the state
    #[inline(always)]
    pub fn refill(&mut self, drounds: u32, out: &mut [u8; BLOCK]) {
        refill_narrow(self, drounds, out)
    }

    #[inline(always)]
    pub(crate) fn refill_rounds(&mut self, drounds: u32) -> State<vec128_storage> {
        refill_narrow_rounds(self, drounds)
    }

    #[inline(always)]
    pub fn set_stream_param(&mut self, param: u32, value: u64) {
        set_stream_param(self, param, value)
    }

    #[inline(always)]
    pub fn get_stream_param(&self, param: u32) -> u64 {
        get_stream_param(self, param)
    }
}

#[inline(always)]
fn refill_wide_impl<Mach: Machine>(
    m: Mach,
    state: &mut ChaCha,
    drounds: u32,
    out: &mut [u8; BUFSZ],
) {
    let k = m.vec([0x6170_7865, 0x3320_646e, 0x7962_2d32, 0x6b20_6574]);
    let mut pos = state.pos64(m);
    let d0: Mach::u32x4 = m.unpack(state.d);
    pos += 1;
    let d1 = d0.insert((pos >> 32) as u32, 1).insert(pos as u32, 0);
    pos += 1;
    let d2 = d0.insert((pos >> 32) as u32, 1).insert(pos as u32, 0);
    pos += 1;
    let d3 = d0.insert((pos >> 32) as u32, 1).insert(pos as u32, 0);

    let b = m.unpack(state.b);
    let c = m.unpack(state.c);
    let mut x = State {
        a: Mach::u32x4x4::from_lanes([k, k, k, k]),
        b: Mach::u32x4x4::from_lanes([b, b, b, b]),
        c: Mach::u32x4x4::from_lanes([c, c, c, c]),
        d: m.unpack(Mach::u32x4x4::from_lanes([d0, d1, d2, d3]).into()),
    };
    for _ in 0..drounds {
        x = round(x);
        x = undiagonalize(round(diagonalize(x)));
    }
    let mut pos = state.pos64(m);
    let d0: Mach::u32x4 = m.unpack(state.d);
    pos += 1;
    let d1 = d0.insert((pos >> 32) as u32, 1).insert(pos as u32, 0);
    pos += 1;
    let d2 = d0.insert((pos >> 32) as u32, 1).insert(pos as u32, 0);
    pos += 1;
    let d3 = d0.insert((pos >> 32) as u32, 1).insert(pos as u32, 0);
    pos += 1;
    let d4 = d0.insert((pos >> 32) as u32, 1).insert(pos as u32, 0);

    let (a, b, c, d) = (
        x.a.to_lanes(),
        x.b.to_lanes(),
        x.c.to_lanes(),
        x.d.to_lanes(),
    );
    let sb = m.unpack(state.b);
    let sc = m.unpack(state.c);
    let sd = [m.unpack(state.d), d1, d2, d3];
    state.d = d4.into();
    let mut words = out.chunks_exact_mut(16);
    for ((((&a, &b), &c), &d), &sd) in a.iter().zip(&b).zip(&c).zip(&d).zip(&sd) {
        (a + k).write_le(words.next().unwrap());
        (b + sb).write_le(words.next().unwrap());
        (c + sc).write_le(words.next().unwrap());
        (d + sd).write_le(words.next().unwrap());
    }
}

dispatch!(m, Mach, {
    fn refill_wide(state: &mut ChaCha, drounds: u32, out: &mut [u8; BUFSZ]) {
        refill_wide_impl(m, state, drounds, out);
    }
});

/// Refill the buffer from a single-block round, updating the block count.
dispatch_light128!(m, Mach, {
    fn refill_narrow(state: &mut ChaCha, drounds: u32, out: &mut [u8; BLOCK]) {
        let x = refill_narrow_rounds(state, drounds);
        let x = State {
            a: m.unpack(x.a),
            b: m.unpack(x.b),
            c: m.unpack(x.c),
            d: m.unpack(x.d),
        };
        state.output_narrow(m, x, out);
        state.inc_block_ct(m);
    }
});

/// Single-block, rounds-only; shared by try_apply_keystream for tails shorter than BUFSZ
/// and XChaCha's setup step.
dispatch!(m, Mach, {
    fn refill_narrow_rounds(state: &mut ChaCha, drounds: u32) -> State<vec128_storage> {
        let k: Mach::u32x4 = m.vec([0x6170_7865, 0x3320_646e, 0x7962_2d32, 0x6b20_6574]);
        let mut x = State {
            a: k,
            b: m.unpack(state.b),
            c: m.unpack(state.c),
            d: m.unpack(state.d),
        };
        for _ in 0..drounds {
            x = round(x);
            x = undiagonalize(round(diagonalize(x)));
        }
        State {
            a: x.a.into(),
            b: x.b.into(),
            c: x.c.into(),
            d: x.d.into(),
        }
    }
});

dispatch_light128!(m, Mach, {
    fn set_stream_param(state: &mut ChaCha, param: u32, value: u64) {
        let d: Mach::u32x4 = m.unpack(state.d);
        state.d = d
            .insert((value >> 32) as u32, (param << 1) | 1)
            .insert(value as u32, param << 1)
            .into();
    }
});

dispatch_light128!(m, Mach, {
    fn get_stream_param(state: &ChaCha, param: u32) -> u64 {
        let d: Mach::u32x4 = m.unpack(state.d);
        ((d.extract((param << 1) | 1) as u64) << 32) | d.extract(param << 1) as u64
    }
});

fn read_u32le(xs: &[u8]) -> u32 {
    assert_eq!(xs.len(), 4);
    u32::from(xs[0]) | (u32::from(xs[1]) << 8) | (u32::from(xs[2]) << 16) | (u32::from(xs[3]) << 24)
}

dispatch_light128!(m, Mach, {
    fn init_chacha(key: &[u8; 32], nonce: &[u8]) -> ChaCha {
        let ctr_nonce = [
            0,
            if nonce.len() == 12 {
                read_u32le(&nonce[0..4])
            } else {
                0
            },
            read_u32le(&nonce[nonce.len() - 8..nonce.len() - 4]),
            read_u32le(&nonce[nonce.len() - 4..]),
        ];
        let key0: Mach::u32x4 = m.read_le(&key[..16]);
        let key1: Mach::u32x4 = m.read_le(&key[16..]);
        ChaCha {
            b: key0.into(),
            c: key1.into(),
            d: ctr_nonce.into(),
        }
    }
});

dispatch_light128!(m, Mach, {
    fn init_chacha_x(key: &[u8; 32], nonce: &[u8; 24], rounds: u32) -> ChaCha {
        let key0: Mach::u32x4 = m.read_le(&key[..16]);
        let key1: Mach::u32x4 = m.read_le(&key[16..]);
        let nonce0: Mach::u32x4 = m.read_le(&nonce[..16]);
        let mut state = ChaCha {
            b: key0.into(),
            c: key1.into(),
            d: nonce0.into(),
        };
        let x = refill_narrow_rounds(&mut state, rounds);
        let ctr_nonce1 = [0, 0, read_u32le(&nonce[16..20]), read_u32le(&nonce[20..24])];
        state.b = x.a;
        state.c = x.d;
        state.d = ctr_nonce1.into();
        state
    }
});
<commit_msg>rand_chacha: remove unused code / eliminate warnings<commit_after>// Copyright 2019 The CryptoCorrosion Contributors
// Copyright 2020 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The ChaCha random number generator.

use ppv_lite86::{dispatch, dispatch_light128};

#[cfg(feature = "rustcrypto_api")]
pub use stream_cipher::generic_array;

pub use ppv_lite86::Machine;
use ppv_lite86::{vec128_storage, ArithOps, BitOps32, LaneWords4, MultiLane, StoreBytes, Vec4};

pub(crate) const BLOCK: usize = 64;
pub(crate) const BLOCK64: u64 = BLOCK as u64;
const LOG2_BUFBLOCKS: u64 = 2;
const BUFBLOCKS: u64 = 1 << LOG2_BUFBLOCKS;
pub(crate) const BUFSZ64: u64 = BLOCK64 * BUFBLOCKS;
pub(crate) const BUFSZ: usize = BUFSZ64 as usize;

#[derive(Clone)]
pub struct ChaCha {
    pub(crate) b: vec128_storage,
    pub(crate) c: vec128_storage,
    pub(crate) d: vec128_storage,
}

#[derive(Clone)]
pub struct State<V> {
    pub(crate) a: V,
    pub(crate) b: V,
    pub(crate) c: V,
    pub(crate) d: V,
}

#[inline(always)]
pub(crate) fn round<V: ArithOps + BitOps32>(mut x: State<V>) -> State<V> {
    x.a += x.b;
    x.d = (x.d ^ x.a).rotate_each_word_right16();
    x.c += x.d;
    x.b = (x.b ^ x.c).rotate_each_word_right20();
    x.a += x.b;
    x.d = (x.d ^ x.a).rotate_each_word_right24();
    x.c += x.d;
    x.b = (x.b ^ x.c).rotate_each_word_right25();
    x
}

#[inline(always)]
pub(crate) fn diagonalize<V: LaneWords4>(mut x: State<V>) -> State<V> {
    x.b = x.b.shuffle_lane_words3012();
    x.c = x.c.shuffle_lane_words2301();
    x.d = x.d.shuffle_lane_words1230();
    x
}
#[inline(always)]
pub(crate) fn undiagonalize<V: LaneWords4>(mut x: State<V>) -> State<V> {
    x.b = x.b.shuffle_lane_words1230();
    x.c = x.c.shuffle_lane_words2301();
    x.d = x.d.shuffle_lane_words3012();
    x
}

impl ChaCha {
    #[inline(always)]
    pub fn new(key: &[u8; 32], nonce: &[u8]) -> Self {
        init_chacha(key, nonce)
    }

    #[inline(always)]
    fn pos64<M: Machine>(&self, m: M) -> u64 {
        let d: M::u32x4 = m.unpack(self.d);
        ((d.extract(1) as u64) << 32) | d.extract(0) as u64
    }

    /// Produce 4 blocks of output, advancing the state
    #[inline(always)]
    pub fn refill4(&mut self, drounds: u32, out: &mut [u8; BUFSZ]) {
        refill_wide(self, drounds, out)
    }

    #[inline(always)]
    pub fn set_stream_param(&mut self, param: u32, value: u64) {
        set_stream_param(self, param, value)
    }

    #[inline(always)]
    pub fn get_stream_param(&self, param: u32) -> u64 {
        get_stream_param(self, param)
    }
}

#[inline(always)]
fn refill_wide_impl<Mach: Machine>(
    m: Mach,
    state: &mut ChaCha,
    drounds: u32,
    out: &mut [u8; BUFSZ],
) {
    let k = m.vec([0x6170_7865, 0x3320_646e, 0x7962_2d32, 0x6b20_6574]);
    let mut pos = state.pos64(m);
    let d0: Mach::u32x4 = m.unpack(state.d);
    pos += 1;
    let d1 = d0.insert((pos >> 32) as u32, 1).insert(pos as u32, 0);
    pos += 1;
    let d2 = d0.insert((pos >> 32) as u32, 1).insert(pos as u32, 0);
    pos += 1;
    let d3 = d0.insert((pos >> 32) as u32, 1).insert(pos as u32, 0);

    let b = m.unpack(state.b);
    let c = m.unpack(state.c);
    let mut x = State {
        a: Mach::u32x4x4::from_lanes([k, k, k, k]),
        b: Mach::u32x4x4::from_lanes([b, b, b, b]),
        c: Mach::u32x4x4::from_lanes([c, c, c, c]),
        d: m.unpack(Mach::u32x4x4::from_lanes([d0, d1, d2, d3]).into()),
    };
    for _ in 0..drounds {
        x = round(x);
        x = undiagonalize(round(diagonalize(x)));
    }
    let mut pos = state.pos64(m);
    let d0: Mach::u32x4 = m.unpack(state.d);
    pos += 1;
    let d1 = d0.insert((pos >> 32) as u32, 1).insert(pos as u32, 0);
    pos += 1;
    let d2 = d0.insert((pos >> 32) as u32, 1).insert(pos as u32, 0);
    pos += 1;
    let d3 = d0.insert((pos >> 32) as u32, 1).insert(pos as u32, 0);
    pos += 1;
    let d4 = d0.insert((pos >> 32) as u32, 1).insert(pos as u32, 0);

    let (a, b, c, d) = (
        x.a.to_lanes(),
        x.b.to_lanes(),
        x.c.to_lanes(),
        x.d.to_lanes(),
    );
    let sb = m.unpack(state.b);
    let sc = m.unpack(state.c);
    let sd = [m.unpack(state.d), d1, d2, d3];
    state.d = d4.into();
    let mut words = out.chunks_exact_mut(16);
    for ((((&a, &b), &c), &d), &sd) in a.iter().zip(&b).zip(&c).zip(&d).zip(&sd) {
        (a + k).write_le(words.next().unwrap());
        (b + sb).write_le(words.next().unwrap());
        (c + sc).write_le(words.next().unwrap());
        (d + sd).write_le(words.next().unwrap());
    }
}

dispatch!(m, Mach, {
    fn refill_wide(state: &mut ChaCha, drounds: u32, out: &mut [u8; BUFSZ]) {
        refill_wide_impl(m, state, drounds, out);
    }
});

// Single-block, rounds-only; shared by try_apply_keystream for tails shorter than BUFSZ
// and XChaCha's setup step.
dispatch!(m, Mach, {
    fn refill_narrow_rounds(state: &mut ChaCha, drounds: u32) -> State<vec128_storage> {
        let k: Mach::u32x4 = m.vec([0x6170_7865, 0x3320_646e, 0x7962_2d32, 0x6b20_6574]);
        let mut x = State {
            a: k,
            b: m.unpack(state.b),
            c: m.unpack(state.c),
            d: m.unpack(state.d),
        };
        for _ in 0..drounds {
            x = round(x);
            x = undiagonalize(round(diagonalize(x)));
        }
        State {
            a: x.a.into(),
            b: x.b.into(),
            c: x.c.into(),
            d: x.d.into(),
        }
    }
});

dispatch_light128!(m, Mach, {
    fn set_stream_param(state: &mut ChaCha, param: u32, value: u64) {
        let d: Mach::u32x4 = m.unpack(state.d);
        state.d = d
            .insert((value >> 32) as u32, (param << 1) | 1)
            .insert(value as u32, param << 1)
            .into();
    }
});

dispatch_light128!(m, Mach, {
    fn get_stream_param(state: &ChaCha, param: u32) -> u64 {
        let d: Mach::u32x4 = m.unpack(state.d);
        ((d.extract((param << 1) | 1) as u64) << 32) | d.extract(param << 1) as u64
    }
});

fn read_u32le(xs: &[u8]) -> u32 {
    assert_eq!(xs.len(), 4);
    u32::from(xs[0]) | (u32::from(xs[1]) << 8) | (u32::from(xs[2]) << 16) | (u32::from(xs[3]) << 24)
}

dispatch_light128!(m, Mach, {
    fn init_chacha(key: &[u8; 32], nonce: &[u8]) -> ChaCha {
        let ctr_nonce = [
            0,
            if nonce.len() == 12 {
                read_u32le(&nonce[0..4])
            } else {
                0
            },
            read_u32le(&nonce[nonce.len() - 8..nonce.len() - 4]),
            read_u32le(&nonce[nonce.len() - 4..]),
        ];
        let key0: Mach::u32x4 = m.read_le(&key[..16]);
        let key1: Mach::u32x4 = m.read_le(&key[16..]);
        ChaCha {
            b: key0.into(),
            c: key1.into(),
            d: ctr_nonce.into(),
        }
    }
});

dispatch_light128!(m, Mach, {
    fn init_chacha_x(key: &[u8; 32], nonce: &[u8; 24], rounds: u32) -> ChaCha {
        let key0: Mach::u32x4 = m.read_le(&key[..16]);
        let key1: Mach::u32x4 = m.read_le(&key[16..]);
        let nonce0: Mach::u32x4 = m.read_le(&nonce[..16]);
        let mut state = ChaCha {
            b: key0.into(),
            c: key1.into(),
            d: nonce0.into(),
        };
        let x = refill_narrow_rounds(&mut state, rounds);
        let ctr_nonce1 = [0, 0, read_u32le(&nonce[16..20]), read_u32le(&nonce[20..24])];
        state.b = x.a;
        state.c = x.d;
        state.d = ctr_nonce1.into();
        state
    }
});
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Fabian Schuiki

//! This crate implements SystemVerilog for the moore compiler.

extern crate llhd;
extern crate moore_common;
pub extern crate moore_svlog_syntax as syntax;
#[macro_use]
extern crate log;

pub(crate) use moore_common as common;

mod ast_map;
pub mod codegen;
mod context;
pub mod hir;

pub use context::*;
pub use syntax::*;
<commit_msg>svlog: add crate prelude<commit_after>// Copyright (c) 2016-2017 Fabian Schuiki

//! This crate implements SystemVerilog for the moore compiler.

extern crate llhd;
extern crate moore_common;
pub extern crate moore_svlog_syntax as syntax;
#[macro_use]
extern crate log;

pub(crate) use moore_common as common;

mod ast_map;
pub mod codegen;
mod context;
pub mod hir;

pub use context::*;
pub use syntax::*;

/// Items commonly used within the crate.
mod crate_prelude {
    pub use ast;
    pub use common::errors::*;
    pub use common::score::Result;
    pub use common::util::{HasDesc, HasSpan};
    pub use common::NodeId;
    pub use context::Context;
}
<|endoftext|>


<commit_before>/// The tokenizer for the query interpreter

use error::*;

#[derive(Debug, PartialEq, Eq)]
pub enum Token {
    Identifier {
        ident: String,
        next: Option<Box<Token>>
    },

    Index {
        idx: usize,
        next: Option<Box<Token>>
    }
}

impl Token {

    pub fn next(&self) -> Option<&Box<Token>> {
        match self {
            &Token::Identifier { ref next, .. } => next.as_ref(),
            &Token::Index { ref next, .. }      => next.as_ref(),
        }
    }

    /// Convenience function for `token.next().is_some()`
    pub fn has_next(&self) -> bool {
        self.next().is_some()
    }

    pub fn set_next(&mut self, token: Token) {
        match self {
            &mut Token::Identifier { ref mut next, .. } => *next = Some(Box::new(token)),
            &mut Token::Index { ref mut next, .. }      => *next = Some(Box::new(token)),
        }
    }

    /// Pop the last token from the chain of tokens
    ///
    /// Returns None if the current Token has no next token
    pub fn pop_last(&mut self) -> Option<Box<Token>> {
        if !self.has_next() {
            None
        } else {
            match self {
                &mut Token::Identifier { next: ref mut next, .. } => {
                    if next.is_some() {
                        let mut n = next.take().unwrap();
                        if n.has_next() {
                            let result = n.pop_last();
                            *next = Some(n);
                            return result;
                        } else {
                            Some(n)
                        }
                    } else {
                        None
                    }
                },

                &mut Token::Index { next: ref mut next, .. } => {
                    if next.is_some() {
                        let mut n = next.take().unwrap();
                        if n.has_next() {
                            let result = n.pop_last();
                            *next = Some(n);
                            return result;
                        } else {
                            Some(n)
                        }
                    } else {
                        None
                    }
                },
            }
        }
    }

    #[cfg(test)]
    pub fn identifier(&self) -> &String {
        match self {
            &Token::Identifier { ref ident, .. } => &ident,
            _ => unreachable!(),
        }
    }

    #[cfg(test)]
    pub fn idx(&self) -> usize {
        match self {
            &Token::Index { idx: i, .. } => i,
            _ => unreachable!(),
        }
    }

}

pub fn tokenize_with_seperator(query: &String, seperator: char) -> Result<Token> {
    use std::str::Split;

    /// Creates a Token object from a string
    ///
    /// # Panics
    ///
    /// * If the internal regex does not compile (should never happen)
    /// * If the token is non-valid (that is, a array index with a non-i64)
    /// * If the regex does not find anything
    /// * If the integer in the brackets (`[]`) cannot be parsed to a valid i64
    ///
    /// # Incorrect behaviour
    ///
    /// * If the regex finds multiple captures
    ///
    /// # Returns
    ///
    /// The `Token` object with the correct identifier/index for this token and no next token.
    ///
    fn mk_token_object(s: &str) -> Result<Token> {
        use regex::Regex;
        use std::str::FromStr;

        lazy_static! {
            static ref RE: Regex = Regex::new(r"^\[\d+\]$").unwrap();
        }

        if !has_array_brackets(s) {
            return Ok(Token::Identifier { ident: String::from(s), next: None });
        }

        match RE.captures(s) {
            None => return Err(Error::from(ErrorKind::ArrayAccessWithoutIndex)),
            Some(captures) => {
                match captures.get(0) {
                    None => Ok(Token::Identifier { ident: String::from(s), next: None }),
                    Some(mtch) => {
                        let mtch = without_array_brackets(mtch.as_str());
                        let i : usize = FromStr::from_str(&mtch).unwrap(); // save because regex
                        Ok(Token::Index {
                            idx: i,
                            next: None,
                        })
                    }
                }
            }
        }
    }

    /// Check whether a str begins with '[' and ends with ']'
    fn has_array_brackets(s: &str) -> bool {
        s.as_bytes()[0] == b'[' && s.as_bytes()[s.len() - 1] == b']'
    }

    /// Remove '[' and ']' from a str
    fn without_array_brackets(s: &str) -> String {
        s.replace("[","").replace("]","")
    }

    fn build_token_tree(split: &mut Split<char>, last: &mut Token) -> Result<()> {
        match split.next() {
            None        => { /* No more tokens */ }
            Some(token) => {
                if token.len() == 0 {
                    return Err(Error::from(ErrorKind::EmptyIdentifier));
                }
                let mut token = try!(mk_token_object(token));
                try!(build_token_tree(split, &mut token));
                last.set_next(token);
            }
        }
        Ok(())
    }

    if query.is_empty() {
        return Err(Error::from(ErrorKind::EmptyQueryError));
    }

    let mut tokens = query.split(seperator);

    match tokens.next() {
        None        => Err(Error::from(ErrorKind::EmptyQueryError)),
        Some(token) => {
            if token.len() == 0 {
                return Err(Error::from(ErrorKind::EmptyIdentifier));
            }
            let mut tok = try!(mk_token_object(token));
            let _       = try!(build_token_tree(&mut tokens, &mut tok));
            Ok(tok)
        }
    }
}

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

    use std::ops::Deref;

    #[test]
    fn test_tokenize_empty_query_to_error() {
        let tokens = tokenize_with_seperator(&String::from(""), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::EmptyQueryError { .. }));
    }

    #[test]
    fn test_tokenize_seperator_only() {
        let tokens = tokenize_with_seperator(&String::from("."), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::EmptyIdentifier { .. }));
    }

    #[test]
    fn test_tokenize_array_brackets_only() {
        let tokens = tokenize_with_seperator(&String::from("[]"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::ArrayAccessWithoutIndex { .. }));
    }

    #[test]
    fn test_tokenize_identifiers_with_array_brackets_only() {
        let tokens = tokenize_with_seperator(&String::from("a.b.c.[]"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::ArrayAccessWithoutIndex { .. }));
    }

    #[test]
    fn test_tokenize_identifiers_in_array_brackets() {
        let tokens = tokenize_with_seperator(&String::from("[a]"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::ArrayAccessWithoutIndex { .. }));
    }

    #[test]
    fn test_tokenize_single_token_query() {
        let tokens = tokenize_with_seperator(&String::from("example"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert!(match tokens {
            Token::Identifier { ref ident, next: None } => {
                assert_eq!("example", ident);
                true
            },
            _ => false,
        });
    }

    #[test]
    fn test_tokenize_double_token_query() {
        let tokens = tokenize_with_seperator(&String::from("a.b"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert!(match tokens {
            Token::Identifier { next: Some(ref next), .. } => {
                assert_eq!("b", next.deref().identifier());
                match next.deref() {
                    &Token::Identifier { next: None, .. } => true,
                    _ => false
                }
            },
            _ => false,
        });
        assert_eq!("a", tokens.identifier());
    }

    #[test]
    fn test_tokenize_ident_then_array_query() {
        let tokens = tokenize_with_seperator(&String::from("a.[0]"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert_eq!("a", tokens.identifier());
        assert!(match tokens {
            Token::Identifier { next: Some(ref next), .. } => {
                match next.deref() {
                    &Token::Index { idx: 0, next: None } => true,
                    _ => false
                }
            },
            _ => false,
        });
    }

    #[test]
    fn test_tokenize_many_idents_then_array_query() {
        let tokens = tokenize_with_seperator(&String::from("a.b.c.[1000]"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert_eq!("a", tokens.identifier());

        let expected =
            Token::Identifier {
                ident: String::from("a"),
                next: Some(Box::new(Token::Identifier {
                    ident: String::from("b"),
                    next: Some(Box::new(Token::Identifier {
                        ident: String::from("c"),
                        next: Some(Box::new(Token::Index {
                            idx: 1000,
                            next: None,
                        })),
                    })),
                })),
            };

        assert_eq!(expected, tokens);
    }

    #[test]
    fn test_tokenize_empty_token_after_good_token() {
        let tokens = tokenize_with_seperator(&String::from("a..b"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::EmptyIdentifier { .. }));
    }

    quickcheck! {
        fn test_array_index(i: usize) -> bool {
            match tokenize_with_seperator(&format!("[{}]", i), '.') {
                Ok(Token::Index { next: None, ..  }) => true,
                _                                    => false,
            }
        }
    }

    #[test]
    fn test_pop_last_token_from_single_identifier_token_is_none() {
        let mut token = Token::Identifier {
            ident: String::from("something"),
            next: None,
        };

        let last = token.pop_last();
        assert!(last.is_none());
    }

    #[test]
    fn test_pop_last_token_from_single_index_token_is_none() {
        let mut token = Token::Index {
            idx: 0,
            next: None,
        };

        let last = token.pop_last();
        assert!(last.is_none());
    }

    #[test]
    fn test_pop_last_token_from_single_identifier_token_is_one() {
        let mut token = Token::Identifier {
            ident: String::from("some"),
            next: Some(Box::new(Token::Identifier {
                ident: String::from("thing"),
                next: None,
            })),
        };

        let last = token.pop_last();

        assert!(last.is_some());
        let last = last.unwrap();

        assert!(is_match!(*last, Token::Identifier { .. }));
        match *last {
            Token::Identifier { ident: ident, .. } => {
                assert_eq!("thing", ident);
            }
            _ => panic!("What just happened?"),
        }
    }

    #[test]
    fn test_pop_last_token_from_single_index_token_is_one() {
        let mut token = Token::Index {
            idx: 0,
            next: Some(Box::new(Token::Index {
                idx: 1,
                next: None,
            })),
        };

        let last = token.pop_last();

        assert!(last.is_some());
        let last = last.unwrap();

        assert!(is_match!(*last, Token::Index { idx: 1, .. }));
    }

    #[test]
    fn test_pop_last_token_from_identifier_chain() {
        let tokens = tokenize_with_seperator(&String::from("a.b.c.d.e.f"), '.');
        assert!(tokens.is_ok());
        let mut tokens = tokens.unwrap();

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!("f", last.unwrap().identifier());
    }

    #[test]
    fn test_pop_last_token_from_mixed_chain() {
        let tokens = tokenize_with_seperator(&String::from("a.[100].c.[3].e.f"), '.');
        assert!(tokens.is_ok());
        let mut tokens = tokens.unwrap();

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!("f", last.unwrap().identifier());
    }

    #[test]
    fn test_pop_last_token_from_identifier_chain_is_array() {
        let tokens = tokenize_with_seperator(&String::from("a.b.c.d.e.f.[1000]"), '.');
        assert!(tokens.is_ok());
        let mut tokens = tokens.unwrap();

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!(1000, last.unwrap().idx());
    }

    #[test]
    fn test_pop_last_token_from_mixed_chain_is_array() {
        let tokens = tokenize_with_seperator(&String::from("a.[100].c.[3].e.f.[1000]"), '.');
        assert!(tokens.is_ok());
        let mut tokens = tokens.unwrap();

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!(1000, last.unwrap().idx());
    }

    #[test]
    fn test_pop_last_token_from_one_token() {
        let tokens = tokenize_with_seperator(&String::from("a"), '.');
        assert!(tokens.is_ok());
        let mut tokens = tokens.unwrap();

        let last = tokens.pop_last();
        assert!(last.is_none());
    }

    #[test]
    fn test_pop_last_chain() {
        let tokens = tokenize_with_seperator(&String::from("a.[100].c.[3].e.f.[1000]"), '.');
        assert!(tokens.is_ok());
        let mut tokens = tokens.unwrap();

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!(1000, last.unwrap().idx());

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!("f", last.unwrap().identifier());

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!("e", last.unwrap().identifier());

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!(3, last.unwrap().idx());

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!("c", last.unwrap().identifier());

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!(100, last.unwrap().idx());

        let last = tokens.pop_last();
        assert!(last.is_none());
    }

}
<commit_msg>Refactor: Remove unneeded names from patterns<commit_after>/// The tokenizer for the query interpreter

use error::*;

#[derive(Debug, PartialEq, Eq)]
pub enum Token {
    Identifier {
        ident: String,
        next: Option<Box<Token>>
    },

    Index {
        idx: usize,
        next: Option<Box<Token>>
    }
}

impl Token {

    pub fn next(&self) -> Option<&Box<Token>> {
        match self {
            &Token::Identifier { ref next, .. } => next.as_ref(),
            &Token::Index { ref next, .. }      => next.as_ref(),
        }
    }

    /// Convenience function for `token.next().is_some()`
    pub fn has_next(&self) -> bool {
        self.next().is_some()
    }

    pub fn set_next(&mut self, token: Token) {
        match self {
            &mut Token::Identifier { ref mut next, .. } => *next = Some(Box::new(token)),
            &mut Token::Index { ref mut next, .. }      => *next = Some(Box::new(token)),
        }
    }

    /// Pop the last token from the chain of tokens
    ///
    /// Returns None if the current Token has no next token
    pub fn pop_last(&mut self) -> Option<Box<Token>> {
        if !self.has_next() {
            None
        } else {
            match self {
                &mut Token::Identifier { ref mut next, .. } => {
                    if next.is_some() {
                        let mut n = next.take().unwrap();
                        if n.has_next() {
                            let result = n.pop_last();
                            *next = Some(n);
                            return result;
                        } else {
                            Some(n)
                        }
                    } else {
                        None
                    }
                },

                &mut Token::Index { ref mut next, .. } => {
                    if next.is_some() {
                        let mut n = next.take().unwrap();
                        if n.has_next() {
                            let result = n.pop_last();
                            *next = Some(n);
                            return result;
                        } else {
                            Some(n)
                        }
                    } else {
                        None
                    }
                },
            }
        }
    }

    #[cfg(test)]
    pub fn identifier(&self) -> &String {
        match self {
            &Token::Identifier { ref ident, .. } => &ident,
            _ => unreachable!(),
        }
    }

    #[cfg(test)]
    pub fn idx(&self) -> usize {
        match self {
            &Token::Index { idx: i, .. } => i,
            _ => unreachable!(),
        }
    }

}

pub fn tokenize_with_seperator(query: &String, seperator: char) -> Result<Token> {
    use std::str::Split;

    /// Creates a Token object from a string
    ///
    /// # Panics
    ///
    /// * If the internal regex does not compile (should never happen)
    /// * If the token is non-valid (that is, a array index with a non-i64)
    /// * If the regex does not find anything
    /// * If the integer in the brackets (`[]`) cannot be parsed to a valid i64
    ///
    /// # Incorrect behaviour
    ///
    /// * If the regex finds multiple captures
    ///
    /// # Returns
    ///
    /// The `Token` object with the correct identifier/index for this token and no next token.
    ///
    fn mk_token_object(s: &str) -> Result<Token> {
        use regex::Regex;
        use std::str::FromStr;

        lazy_static! {
            static ref RE: Regex = Regex::new(r"^\[\d+\]$").unwrap();
        }

        if !has_array_brackets(s) {
            return Ok(Token::Identifier { ident: String::from(s), next: None });
        }

        match RE.captures(s) {
            None => return Err(Error::from(ErrorKind::ArrayAccessWithoutIndex)),
            Some(captures) => {
                match captures.get(0) {
                    None => Ok(Token::Identifier { ident: String::from(s), next: None }),
                    Some(mtch) => {
                        let mtch = without_array_brackets(mtch.as_str());
                        let i : usize = FromStr::from_str(&mtch).unwrap(); // save because regex
                        Ok(Token::Index {
                            idx: i,
                            next: None,
                        })
                    }
                }
            }
        }
    }

    /// Check whether a str begins with '[' and ends with ']'
    fn has_array_brackets(s: &str) -> bool {
        s.as_bytes()[0] == b'[' && s.as_bytes()[s.len() - 1] == b']'
    }

    /// Remove '[' and ']' from a str
    fn without_array_brackets(s: &str) -> String {
        s.replace("[","").replace("]","")
    }

    fn build_token_tree(split: &mut Split<char>, last: &mut Token) -> Result<()> {
        match split.next() {
            None        => { /* No more tokens */ }
            Some(token) => {
                if token.len() == 0 {
                    return Err(Error::from(ErrorKind::EmptyIdentifier));
                }
                let mut token = try!(mk_token_object(token));
                try!(build_token_tree(split, &mut token));
                last.set_next(token);
            }
        }
        Ok(())
    }

    if query.is_empty() {
        return Err(Error::from(ErrorKind::EmptyQueryError));
    }

    let mut tokens = query.split(seperator);

    match tokens.next() {
        None        => Err(Error::from(ErrorKind::EmptyQueryError)),
        Some(token) => {
            if token.len() == 0 {
                return Err(Error::from(ErrorKind::EmptyIdentifier));
            }
            let mut tok = try!(mk_token_object(token));
            let _       = try!(build_token_tree(&mut tokens, &mut tok));
            Ok(tok)
        }
    }
}

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

    use std::ops::Deref;

    #[test]
    fn test_tokenize_empty_query_to_error() {
        let tokens = tokenize_with_seperator(&String::from(""), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::EmptyQueryError { .. }));
    }

    #[test]
    fn test_tokenize_seperator_only() {
        let tokens = tokenize_with_seperator(&String::from("."), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::EmptyIdentifier { .. }));
    }

    #[test]
    fn test_tokenize_array_brackets_only() {
        let tokens = tokenize_with_seperator(&String::from("[]"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::ArrayAccessWithoutIndex { .. }));
    }

    #[test]
    fn test_tokenize_identifiers_with_array_brackets_only() {
        let tokens = tokenize_with_seperator(&String::from("a.b.c.[]"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::ArrayAccessWithoutIndex { .. }));
    }

    #[test]
    fn test_tokenize_identifiers_in_array_brackets() {
        let tokens = tokenize_with_seperator(&String::from("[a]"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::ArrayAccessWithoutIndex { .. }));
    }

    #[test]
    fn test_tokenize_single_token_query() {
        let tokens = tokenize_with_seperator(&String::from("example"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert!(match tokens {
            Token::Identifier { ref ident, next: None } => {
                assert_eq!("example", ident);
                true
            },
            _ => false,
        });
    }

    #[test]
    fn test_tokenize_double_token_query() {
        let tokens = tokenize_with_seperator(&String::from("a.b"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert!(match tokens {
            Token::Identifier { next: Some(ref next), .. } => {
                assert_eq!("b", next.deref().identifier());
                match next.deref() {
                    &Token::Identifier { next: None, .. } => true,
                    _ => false
                }
            },
            _ => false,
        });
        assert_eq!("a", tokens.identifier());
    }

    #[test]
    fn test_tokenize_ident_then_array_query() {
        let tokens = tokenize_with_seperator(&String::from("a.[0]"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert_eq!("a", tokens.identifier());
        assert!(match tokens {
            Token::Identifier { next: Some(ref next), .. } => {
                match next.deref() {
                    &Token::Index { idx: 0, next: None } => true,
                    _ => false
                }
            },
            _ => false,
        });
    }

    #[test]
    fn test_tokenize_many_idents_then_array_query() {
        let tokens = tokenize_with_seperator(&String::from("a.b.c.[1000]"), '.');
        assert!(tokens.is_ok());
        let tokens = tokens.unwrap();

        assert_eq!("a", tokens.identifier());

        let expected =
            Token::Identifier {
                ident: String::from("a"),
                next: Some(Box::new(Token::Identifier {
                    ident: String::from("b"),
                    next: Some(Box::new(Token::Identifier {
                        ident: String::from("c"),
                        next: Some(Box::new(Token::Index {
                            idx: 1000,
                            next: None,
                        })),
                    })),
                })),
            };

        assert_eq!(expected, tokens);
    }

    #[test]
    fn test_tokenize_empty_token_after_good_token() {
        let tokens = tokenize_with_seperator(&String::from("a..b"), '.');
        assert!(tokens.is_err());
        let tokens = tokens.unwrap_err();

        let errkind = tokens.kind();
        assert!(is_match!(errkind, &ErrorKind::EmptyIdentifier { .. }));
    }

    quickcheck! {
        fn test_array_index(i: usize) -> bool {
            match tokenize_with_seperator(&format!("[{}]", i), '.') {
                Ok(Token::Index { next: None, ..  }) => true,
                _                                    => false,
            }
        }
    }

    #[test]
    fn test_pop_last_token_from_single_identifier_token_is_none() {
        let mut token = Token::Identifier {
            ident: String::from("something"),
            next: None,
        };

        let last = token.pop_last();
        assert!(last.is_none());
    }

    #[test]
    fn test_pop_last_token_from_single_index_token_is_none() {
        let mut token = Token::Index {
            idx: 0,
            next: None,
        };

        let last = token.pop_last();
        assert!(last.is_none());
    }

    #[test]
    fn test_pop_last_token_from_single_identifier_token_is_one() {
        let mut token = Token::Identifier {
            ident: String::from("some"),
            next: Some(Box::new(Token::Identifier {
                ident: String::from("thing"),
                next: None,
            })),
        };

        let last = token.pop_last();

        assert!(last.is_some());
        let last = last.unwrap();

        assert!(is_match!(*last, Token::Identifier { .. }));
        match *last {
            Token::Identifier { ident, .. } => {
                assert_eq!("thing", ident);
            }
            _ => panic!("What just happened?"),
        }
    }

    #[test]
    fn test_pop_last_token_from_single_index_token_is_one() {
        let mut token = Token::Index {
            idx: 0,
            next: Some(Box::new(Token::Index {
                idx: 1,
                next: None,
            })),
        };

        let last = token.pop_last();

        assert!(last.is_some());
        let last = last.unwrap();

        assert!(is_match!(*last, Token::Index { idx: 1, .. }));
    }

    #[test]
    fn test_pop_last_token_from_identifier_chain() {
        let tokens = tokenize_with_seperator(&String::from("a.b.c.d.e.f"), '.');
        assert!(tokens.is_ok());
        let mut tokens = tokens.unwrap();

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!("f", last.unwrap().identifier());
    }

    #[test]
    fn test_pop_last_token_from_mixed_chain() {
        let tokens = tokenize_with_seperator(&String::from("a.[100].c.[3].e.f"), '.');
        assert!(tokens.is_ok());
        let mut tokens = tokens.unwrap();

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!("f", last.unwrap().identifier());
    }

    #[test]
    fn test_pop_last_token_from_identifier_chain_is_array() {
        let tokens = tokenize_with_seperator(&String::from("a.b.c.d.e.f.[1000]"), '.');
        assert!(tokens.is_ok());
        let mut tokens = tokens.unwrap();

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!(1000, last.unwrap().idx());
    }

    #[test]
    fn test_pop_last_token_from_mixed_chain_is_array() {
        let tokens = tokenize_with_seperator(&String::from("a.[100].c.[3].e.f.[1000]"), '.');
        assert!(tokens.is_ok());
        let mut tokens = tokens.unwrap();

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!(1000, last.unwrap().idx());
    }

    #[test]
    fn test_pop_last_token_from_one_token() {
        let tokens = tokenize_with_seperator(&String::from("a"), '.');
        assert!(tokens.is_ok());
        let mut tokens = tokens.unwrap();

        let last = tokens.pop_last();
        assert!(last.is_none());
    }

    #[test]
    fn test_pop_last_chain() {
        let tokens = tokenize_with_seperator(&String::from("a.[100].c.[3].e.f.[1000]"), '.');
        assert!(tokens.is_ok());
        let mut tokens = tokens.unwrap();

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!(1000, last.unwrap().idx());

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!("f", last.unwrap().identifier());

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!("e", last.unwrap().identifier());

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!(3, last.unwrap().idx());

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!("c", last.unwrap().identifier());

        let last = tokens.pop_last();
        assert!(last.is_some());
        assert_eq!(100, last.unwrap().idx());

        let last = tokens.pop_last();
        assert!(last.is_none());
    }

}
<|endoftext|>