<commit_before>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the
// terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not
// distributed with this file, You can
// obtain one at
// http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::{str, ptr};
use std::ffi::{CStr, CString};
use std::thread;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender, Receiver};

use super::libc::{c_int, c_char};
use super::simple_stream::bstream::Bstream;

extern "C" {
    fn register_writer_tx(tx: *mut Sender<Vec<u8>>);
    fn register_stop_tx(tx: *mut Sender<()>);
}

#[no_mangle]
pub extern "C" fn some_new_fn() -> c_int {
    println!("some_new_fn");
    65 as c_int
}

#[no_mangle]
pub extern "C" fn start(address: *const c_char,
    handler: extern fn(*const c_char, c_int),
    on_connect_handler: extern fn(),
    on_disconnect_handler: extern fn()) -> c_int {

    println!("start");
    88 as c_int
}

/// Connects to the provided address, (eg "123.123.123.123:3000") and registers
/// the on data received handler
#[no_mangle]
pub extern "C" fn connect(address: *const c_char,
    handler: extern fn(*const c_char, c_int),
    on_connect_handler: extern fn(),
    on_disconnect_handler: extern fn()) -> c_int {

    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");

    78 as c_int

    // let mut r_address;
    // unsafe {
    //     r_address = CStr::from_ptr(address);
    // }
    // let s_address = r_address.to_bytes();
    // let host_address = match str::from_utf8(s_address) {
    //     Ok(safe_str) => safe_str,
    //     Err(_) => {
    //         println!("Invalid host address");
    //         return -1 as c_int;
    //     }
    // };
    //
    // // Create and register a way to kill this client
    // let (k_tx, kill_rx): (Sender<()>, Receiver<()>) = channel();
    // let kill_tx = k_tx.clone();
    // let mut k_tx_ptr = Box::new(k_tx);
    // unsafe {
    //     register_stop_tx(&mut *k_tx_ptr);
    // }
    //
    // // Writer thread's channel
    // let (w_tx, w_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
    // let mut w_tx_ptr = Box::new(w_tx);
    // unsafe {
    //     register_writer_tx(&mut *w_tx_ptr);
    // }
    //
    // println!("Attempting connect to: {}", host_address);
    //
    // let result = TcpStream::connect(host_address);
    // if result.is_err() {
    //     println!("Error connecting to {} - {}", host_address, result.unwrap_err());
    //     return -1 as c_int;
    // }
    // println!("Connected");
    // on_connect_handler();
    //
    // let stream = result.unwrap();
    // let client = Bstream::new(stream);
    //
    // let r_client = client.clone();
    // let r_kill_tx = kill_tx.clone();
    //
    // let w_client = client.clone();
    // let w_kill_tx = kill_tx.clone();
    //
    // // Start the reader thread
    // thread::Builder::new()
    //     .name("ReaderThread".to_string())
    //     .spawn(move||{
    //         reader_thread(r_client, handler, r_kill_tx)
    //     }).unwrap();
    //
    // // Start the writer thread
    // thread::Builder::new()
    //     .name("WriterThread".to_string())
    //     .spawn(move||{
    //         writer_thread(w_rx, w_client, w_kill_tx)
    //     }).unwrap();
    //
    // // Wait for the kill signal
    // match kill_rx.recv() {
    //     Ok(_) => { }
    //     Err(e) => {
    //         println!("Error on kill channel: {}", e);
    //         return -1 as c_int;
    //     }
    // };
    // on_disconnect_handler();
    //
    // // Exit out in standard C fashion
    // 0 as c_int
}

/// Writes the complete contents of buffer to the server
/// Returns -1 on error
#[no_mangle]
pub extern "C" fn send_to_writer(w_tx: *mut Sender<Vec<u8>>,
                                 buffer: *const c_char,
                                 count: c_int,
                                 k_tx: *mut Sender<()>) -> c_int {
    if count < 1 {
        println!("Error - count must be greater than zero");
        return -1 as c_int;
    }

    let num_elts = count as usize;
    let mut n_buffer = Vec::<u8>::with_capacity(num_elts);
    for x in 0..num_elts as isize {
        unsafe {
            n_buffer.push(ptr::read(buffer.offset(x)) as u8);
        }
    }

    unsafe {
        match (*w_tx).send(n_buffer) {
            Ok(_) => { }
            Err(e) => {
                println!("Error sending buffer: {}", e);
                let _ = (*k_tx).send(());
                return -1 as c_int;
            }
        };
    }

    0 as c_int
}

/// Forever listens to incoming data and when a complete message is received,
/// the passed callback is hit
fn reader_thread(client: Bstream,
                 event_handler: extern fn(*const c_char, c_int),
                 kill_tx: Sender<()>) {

    let mut reader = client.clone();
    loop {
        match reader.read() {
            Ok(buffer) => {
                // Launch the handler in a new thread
                thread::Builder::new()
                    .name("Reader-Worker".to_string())
                    .spawn(move||{
                        let slice = &buffer[..];
                        let c_buffer = CString::new(slice).unwrap();
                        event_handler(c_buffer.as_ptr(), buffer.len() as c_int);
                    }).unwrap();
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Reader thread finished");
    let _ = kill_tx.send(());
}

/// Forever listens to Receiver<Vec<u8>> waiting on messages to come in
/// Once available, blocks until the entire message has been written
fn writer_thread(rx: Receiver<Vec<u8>>, client: Bstream, kill_tx: Sender<()>) {
    let mut writer = client.clone();
    loop {
        match rx.recv() {
            Ok(ref mut buffer) => {
                match writer.write(buffer) {
                    Ok(_) => { }
                    Err(e) => {
                        println!("Error: {}", e);
                        break;
                    }
                };
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Writer thread finished");
    let _ = kill_tx.send(());
}
<commit_msg>Debugging<commit_after>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the
// terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not
// distributed with this file, You can
// obtain one at
// http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::{str, ptr};
use std::ffi::{CStr, CString};
use std::thread;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender, Receiver};

use super::libc::{c_int, c_char};
use super::simple_stream::bstream::Bstream;

extern "C" {
    fn register_writer_tx(tx: *mut Sender<Vec<u8>>);
    fn register_stop_tx(tx: *mut Sender<()>);
}

// #[no_mangle]
// pub extern "C" fn some_new_fn() -> c_int {
//     println!("some_new_fn");
//     65 as c_int
// }

#[no_mangle]
pub extern "C" fn start(address: *const c_char,
    handler: extern fn(*const c_char, c_int),
    on_connect_handler: extern fn(),
    on_disconnect_handler: extern fn()) -> c_int {

    // println!("start");
    // 88 as c_int

    let mut r_address;
    unsafe {
        r_address = CStr::from_ptr(address);
    }
    let s_address = r_address.to_bytes();
    let host_address = match str::from_utf8(s_address) {
        Ok(safe_str) => safe_str,
        Err(_) => {
            println!("Invalid host address");
            return -1 as c_int;
        }
    };

    // Create and register a way to kill this client
    let (k_tx, kill_rx): (Sender<()>, Receiver<()>) = channel();
    let kill_tx = k_tx.clone();
    let mut k_tx_ptr = Box::new(k_tx);
    unsafe {
        register_stop_tx(&mut *k_tx_ptr);
    }

    // Writer thread's channel
    let (w_tx, w_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
    let mut w_tx_ptr = Box::new(w_tx);
    unsafe {
        register_writer_tx(&mut *w_tx_ptr);
    }

    println!("Attempting connect to: {}", host_address);

    let result = TcpStream::connect(host_address);
    if result.is_err() {
        println!("Error connecting to {} - {}", host_address, result.unwrap_err());
        return -1 as c_int;
    }
    println!("Connected");
    on_connect_handler();

    let stream = result.unwrap();
    let client = Bstream::new(stream);

    let r_client = client.clone();
    let r_kill_tx = kill_tx.clone();

    let w_client = client.clone();
    let w_kill_tx = kill_tx.clone();

    // Start the reader thread
    thread::Builder::new()
        .name("ReaderThread".to_string())
        .spawn(move||{
            reader_thread(r_client, handler, r_kill_tx)
        }).unwrap();

    // Start the writer thread
    thread::Builder::new()
        .name("WriterThread".to_string())
        .spawn(move||{
            writer_thread(w_rx, w_client, w_kill_tx)
        }).unwrap();

    // Wait for the kill signal
    match kill_rx.recv() {
        Ok(_) => { }
        Err(e) => {
            println!("Error on kill channel: {}", e);
            return -1 as c_int;
        }
    };
    on_disconnect_handler();

    // Exit out in standard C fashion
    0 as c_int
}

// /// Connects to the provided address, (eg "123.123.123.123:3000") and registers
// /// the on data received handler
// #[no_mangle]
// pub extern "C" fn connect(address: *const c_char,
//     handler: extern fn(*const c_char, c_int),
//     on_connect_handler: extern fn(),
//     on_disconnect_handler: extern fn()) -> c_int {
//
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//     println!("Rust - connect");
//
//
//     78 as c_int
//
//
//     let mut r_address;
//     unsafe {
//         r_address = CStr::from_ptr(address);
//     }
//     let s_address = r_address.to_bytes();
//     let host_address = match str::from_utf8(s_address) {
//         Ok(safe_str) => safe_str,
//         Err(_) => {
//             println!("Invalid host address");
//             return -1 as c_int;
//         }
//     };
//
//     // Create and register a way to kill this client
//     let (k_tx, kill_rx): (Sender<()>, Receiver<()>) = channel();
//     let kill_tx = k_tx.clone();
//     let mut k_tx_ptr = Box::new(k_tx);
//     unsafe {
//         register_stop_tx(&mut *k_tx_ptr);
//     }
//
//     // Writer thread's channel
//     let (w_tx, w_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
//     let mut w_tx_ptr = Box::new(w_tx);
//     unsafe {
//         register_writer_tx(&mut *w_tx_ptr);
//     }
//
//     println!("Attempting connect to: {}", host_address);
//
//     let result = TcpStream::connect(host_address);
//     if result.is_err() {
//         println!("Error connecting to {} - {}", host_address, result.unwrap_err());
//         return -1 as c_int;
//     }
//     println!("Connected");
//     on_connect_handler();
//
//     let stream = result.unwrap();
//     let client = Bstream::new(stream);
//
//     let r_client = client.clone();
//     let r_kill_tx = kill_tx.clone();
//
//     let w_client = client.clone();
//     let w_kill_tx = kill_tx.clone();
//
//     // Start the reader thread
//     thread::Builder::new()
//         .name("ReaderThread".to_string())
//         .spawn(move||{
//             reader_thread(r_client, handler, r_kill_tx)
//         }).unwrap();
//
//     // Start the writer thread
//     thread::Builder::new()
//         .name("WriterThread".to_string())
//         .spawn(move||{
//             writer_thread(w_rx, w_client, w_kill_tx)
//         }).unwrap();
//
//     // Wait for the kill signal
//     match kill_rx.recv() {
//         Ok(_) => { }
//         Err(e) => {
//             println!("Error on kill channel: {}", e);
//             return -1 as c_int;
//         }
//     };
//     on_disconnect_handler();
//
//     // Exit out in standard C fashion
//     0 as c_int
// }

/// Writes the complete contents of buffer to the server
/// Returns -1 on error
#[no_mangle]
pub extern "C" fn send_to_writer(w_tx: *mut Sender<Vec<u8>>,
                                 buffer: *const c_char,
                                 count: c_int,
                                 k_tx: *mut Sender<()>) -> c_int {
    if count < 1 {
        println!("Error - count must be greater than zero");
        return -1 as c_int;
    }

    let num_elts = count as usize;
    let mut n_buffer = Vec::<u8>::with_capacity(num_elts);
    for x in 0..num_elts as isize {
        unsafe {
            n_buffer.push(ptr::read(buffer.offset(x)) as u8);
        }
    }

    unsafe {
        match (*w_tx).send(n_buffer) {
            Ok(_) => { }
            Err(e) => {
                println!("Error sending buffer: {}", e);
                let _ = (*k_tx).send(());
                return -1 as c_int;
            }
        };
    }

    0 as c_int
}

/// Forever listens to incoming data and when a complete message is received,
/// the passed callback is hit
fn reader_thread(client: Bstream,
                 event_handler: extern fn(*const c_char, c_int),
                 kill_tx: Sender<()>) {

    let mut reader = client.clone();
    loop {
        match reader.read() {
            Ok(buffer) => {
                // Launch the handler in a new thread
                thread::Builder::new()
                    .name("Reader-Worker".to_string())
                    .spawn(move||{
                        let slice = &buffer[..];
                        let c_buffer = CString::new(slice).unwrap();
                        event_handler(c_buffer.as_ptr(), buffer.len() as c_int);
                    }).unwrap();
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Reader thread finished");
    let _ = kill_tx.send(());
}

/// Forever listens to Receiver<Vec<u8>> waiting on messages to come in
/// Once available, blocks until the entire message has been written
fn writer_thread(rx: Receiver<Vec<u8>>, client: Bstream, kill_tx: Sender<()>) {
    let mut writer = client.clone();
    loop {
        match rx.recv() {
            Ok(ref mut buffer) => {
                match writer.write(buffer) {
                    Ok(_) => { }
                    Err(e) => {
                        println!("Error: {}", e);
                        break;
                    }
                };
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Writer thread finished");
    let _ = kill_tx.send(());
}
<|endoftext|>


<commit_before>use std::sync::Arc;
use std::io;
use bifrost::conshash::{ConsistentHashing, CHError};
use bifrost::raft::client::{RaftClient, ClientError};
use bifrost::raft;
use bifrost::rpc::{RPCError, DEFAULT_CLIENT_POOL, Server as RPCServer};
use bifrost::raft::state_machine::master::ExecError;
use bifrost::raft::state_machine::callback::server::{NotifyError};

use server::{transactions as txn_server, cell_rpc as plain_server};
use ram::types::Id;
use ram::cell::{Cell, Header, ReadError, WriteError};
use ram::schema::{sm as schema_sm};
use ram::schema::sm::client::{SMClient as SchemaClient};
use ram::schema::Schema;
use self::transaction::*;

static TRANSACTION_MAX_RETRY: u32 = 100;

pub mod transaction;

#[derive(Debug)]
pub enum NebClientError {
    RaftClientError(ClientError),
    ConsistentHashtableError(CHError)
}

pub struct Client {
    pub conshash: Arc<ConsistentHashing>,
    pub raft_client: Arc<RaftClient>,
    pub schema_client: SchemaClient
}

impl Client {
    pub fn new<'a>(subscription_server: &Arc<RPCServer>, meta_servers: &Vec<String>, group: &'a str) -> Result<Client, NebClientError> {
        match RaftClient::new(meta_servers, raft::DEFAULT_SERVICE_ID) {
            Ok(raft_client) => {
                RaftClient::prepare_subscription(subscription_server);
                assert!(RaftClient::can_callback());
                match ConsistentHashing::new_client(group, &raft_client) {
                    Ok(chash) => Ok(Client {
                        conshash: chash,
                        raft_client: raft_client.clone(),
                        schema_client: SchemaClient::new(schema_sm::generate_sm_id(group), &raft_client)
                    }),
                    Err(err) => Err(NebClientError::ConsistentHashtableError(err))
                }
            },
            Err(err) => Err(NebClientError::RaftClientError(err))
        }
    }
    pub fn locate_server_address(&self, id: &Id) -> Result<String, RPCError> {
        match self.conshash.get_server(id.higher) {
            Some(n) => Ok(n),
            None => Err(RPCError::IOError(io::Error::new(io::ErrorKind::NotFound, "cannot locate")))
        }
    }
    pub fn locate_plain_server(&self, id: &Id) -> Result<Arc<plain_server::SyncServiceClient>, RPCError> {
        let address = self.locate_server_address(id)?;
        let client = match DEFAULT_CLIENT_POOL.get(&address) {
            Ok(c) => c,
            Err(e) => return Err(RPCError::IOError(e))
        };
        Ok(plain_server::SyncServiceClient::new(plain_server::DEFAULT_SERVICE_ID, &client))
    }
    pub fn read_cell(&self, id: &Id) -> Result<Result<Cell, ReadError>, RPCError> {
        let client = self.locate_plain_server(id)?;
        client.read_cell(id)
    }
    pub fn write_cell(&self, cell: &Cell) -> Result<Result<Header, WriteError>, RPCError> {
        let client = self.locate_plain_server(&cell.id())?;
        client.write_cell(cell)
    }
    pub fn update_cell(&self, cell: &Cell) -> Result<Result<Header, WriteError>, RPCError> {
        let client = self.locate_plain_server(&cell.id())?;
        client.update_cell(cell)
    }
    pub fn remove_cell(&self, id: &Id) -> Result<Result<(), WriteError>, RPCError> {
        let client = self.locate_plain_server(id)?;
        client.remove_cell(id)
    }
    pub fn transaction<TFN>(&self, func: TFN) -> Result<(), TxnError>
        where TFN: Fn(&mut Transaction) -> Result<(), TxnError> {
        let server_name = match self.conshash.rand_server() {
            Some(name) => name,
            None => return Err(TxnError::CannotFindAServer)
        };
        let txn_client = match txn_server::new_client(&server_name) {
            Ok(client) => client,
            Err(e) => return Err(TxnError::IoError(e))
        };
        let mut txn_id: txn_server::TxnId;
        let mut retried = 0;
        while retried < TRANSACTION_MAX_RETRY {
            txn_id = match txn_client.begin() {
                Ok(Ok(id)) => id,
                _ => return Err(TxnError::CannotBegin)
            };
            let mut txn = Transaction{
                tid: txn_id,
                state: txn_server::TxnState::Started,
                client: txn_client.clone(),
                changes: 0
            };
            let mut exec_result = func(&mut txn);
            if exec_result.is_ok() && txn.state == txn_server::TxnState::Started {
                exec_result = txn.prepare();
            }
            if exec_result.is_ok() && txn.state == txn_server::TxnState::Prepared {
                exec_result = txn.commit();
            }
            match exec_result {
                Ok(()) => {
                    return Ok(());
                },
                Err(TxnError::NotRealizable) => {
                    txn.abort();  // continue the loop to retry
                },
                Err(e) => {
                    // abort will always be an error to achieve early break
                    txn.abort();
                    return Err(e);
                }
            }
            retried += 1;
            debug!("Client retry transaction, {:?} times", retried);
        }
        Err(TxnError::TooManyRetry)
    }
    pub fn new_schema_with_id(&self, schema: &Schema) -> Result<Result<(), NotifyError>, ExecError> {
        self.schema_client.new_schema(schema)
    }
    pub fn new_schema(&self, schema: &mut Schema) -> Result<Result<(), NotifyError>, ExecError> {
        let schema_id = self.schema_client.next_id()?.unwrap();
        schema.id = schema_id;
        self.new_schema_with_id(schema)
    }
    pub fn del_schema(&self, schema_id: &String) -> Result<Result<(), NotifyError>, ExecError> {
        self.schema_client.del_schema(schema_id)
    }
    pub fn get_all_schema(&self) -> Result<Vec<Schema>, ExecError> {
        Ok(self.schema_client.get_all()?.unwrap())
    }
}<commit_msg>make transactrion return final result value<commit_after>use std::sync::Arc;
use std::io;
use bifrost::conshash::{ConsistentHashing, CHError};
use bifrost::raft::client::{RaftClient, ClientError};
use bifrost::raft;
use bifrost::rpc::{RPCError, DEFAULT_CLIENT_POOL, Server as RPCServer};
use bifrost::raft::state_machine::master::ExecError;
use bifrost::raft::state_machine::callback::server::{NotifyError};

use server::{transactions as txn_server, cell_rpc as plain_server};
use ram::types::Id;
use ram::cell::{Cell, Header, ReadError, WriteError};
use ram::schema::{sm as schema_sm};
use ram::schema::sm::client::{SMClient as SchemaClient};
use ram::schema::Schema;
use self::transaction::*;

static TRANSACTION_MAX_RETRY: u32 = 100;

pub mod transaction;

#[derive(Debug)]
pub enum NebClientError {
    RaftClientError(ClientError),
    ConsistentHashtableError(CHError)
}

pub struct Client {
    pub conshash: Arc<ConsistentHashing>,
    pub raft_client: Arc<RaftClient>,
    pub schema_client: SchemaClient
}

impl Client {
    pub fn new<'a>(subscription_server: &Arc<RPCServer>, meta_servers: &Vec<String>, group: &'a str) -> Result<Client, NebClientError> {
        match RaftClient::new(meta_servers, raft::DEFAULT_SERVICE_ID) {
            Ok(raft_client) => {
                RaftClient::prepare_subscription(subscription_server);
                assert!(RaftClient::can_callback());
                match ConsistentHashing::new_client(group, &raft_client) {
                    Ok(chash) => Ok(Client {
                        conshash: chash,
                        raft_client: raft_client.clone(),
                        schema_client: SchemaClient::new(schema_sm::generate_sm_id(group), &raft_client)
                    }),
                    Err(err) => Err(NebClientError::ConsistentHashtableError(err))
                }
            },
            Err(err) => Err(NebClientError::RaftClientError(err))
        }
    }
    pub fn locate_server_address(&self, id: &Id) -> Result<String, RPCError> {
        match self.conshash.get_server(id.higher) {
            Some(n) => Ok(n),
            None => Err(RPCError::IOError(io::Error::new(io::ErrorKind::NotFound, "cannot locate")))
        }
    }
    pub fn locate_plain_server(&self, id: &Id) -> Result<Arc<plain_server::SyncServiceClient>, RPCError> {
        let address = self.locate_server_address(id)?;
        let client = match DEFAULT_CLIENT_POOL.get(&address) {
            Ok(c) => c,
            Err(e) => return Err(RPCError::IOError(e))
        };
        Ok(plain_server::SyncServiceClient::new(plain_server::DEFAULT_SERVICE_ID, &client))
    }
    pub fn read_cell(&self, id: &Id) -> Result<Result<Cell, ReadError>, RPCError> {
        let client = self.locate_plain_server(id)?;
        client.read_cell(id)
    }
    pub fn write_cell(&self, cell: &Cell) -> Result<Result<Header, WriteError>, RPCError> {
        let client = self.locate_plain_server(&cell.id())?;
        client.write_cell(cell)
    }
    pub fn update_cell(&self, cell: &Cell) -> Result<Result<Header, WriteError>, RPCError> {
        let client = self.locate_plain_server(&cell.id())?;
        client.update_cell(cell)
    }
    pub fn remove_cell(&self, id: &Id) -> Result<Result<(), WriteError>, RPCError> {
        let client = self.locate_plain_server(id)?;
        client.remove_cell(id)
    }
    pub fn transaction<TFN, TR>(&self, func: TFN) -> Result<TR, TxnError>
        where TFN: Fn(&mut Transaction) -> Result<TR, TxnError> {
        let server_name = match self.conshash.rand_server() {
            Some(name) => name,
            None => return Err(TxnError::CannotFindAServer)
        };
        let txn_client = match txn_server::new_client(&server_name) {
            Ok(client) => client,
            Err(e) => return Err(TxnError::IoError(e))
        };
        let mut txn_id: txn_server::TxnId;
        let mut retried = 0;
        while retried < TRANSACTION_MAX_RETRY {
            txn_id = match txn_client.begin() {
                Ok(Ok(id)) => id,
                _ => return Err(TxnError::CannotBegin)
            };
            let mut txn = Transaction{
                tid: txn_id,
                state: txn_server::TxnState::Started,
                client: txn_client.clone(),
                changes: 0
            };
            let exec_result = func(&mut txn);
            let mut exec_value = None;
            let mut txn_result = Ok(());
            match exec_result {
                Ok(val) => {
                    if txn.state == txn_server::TxnState::Started {
                        txn_result = txn.prepare();
                    }
                    if txn_result.is_ok() && txn.state == txn_server::TxnState::Prepared {
                        txn_result = txn.commit();
                    }
                    exec_value = Some(val);
                },
                Err(e) => {
                    txn_result = Err(e)
                }
            }
            match txn_result {
                Ok(()) => {
                    return Ok(exec_value.unwrap());
                },
                Err(TxnError::NotRealizable) => {
                    txn.abort();  // continue the loop to retry
                },
                Err(e) => {
                    // abort will always be an error to achieve early break
                    txn.abort();
                    return Err(e);
                }
            }
            retried += 1;
            debug!("Client retry transaction, {:?} times", retried);
        }
        Err(TxnError::TooManyRetry)
    }
    pub fn new_schema_with_id(&self, schema: &Schema) -> Result<Result<(), NotifyError>, ExecError> {
        self.schema_client.new_schema(schema)
    }
    pub fn new_schema(&self, schema: &mut Schema) -> Result<Result<(), NotifyError>, ExecError> {
        let schema_id = self.schema_client.next_id()?.unwrap();
        schema.id = schema_id;
        self.new_schema_with_id(schema)
    }
    pub fn del_schema(&self, schema_id: &String) -> Result<Result<(), NotifyError>, ExecError> {
        self.schema_client.del_schema(schema_id)
    }
    pub fn get_all_schema(&self) -> Result<Vec<Schema>, ExecError> {
        Ok(self.schema_client.get_all()?.unwrap())
    }
}<|endoftext|>


<commit_before>#![cfg(feature = "encoding")]

#![feature(test)]
extern crate test;

#[macro_use]
extern crate lambda_calculus as lambda;

use test::Bencher;
use lambda::church::list::*;
use lambda::church::numerals::*;
use lambda::church::boolean::fls;
use lambda::combinators::C;
use lambda::*;

fn list1() -> Term { vec![1].into_church() }
fn list2() -> Term { vec![2, 3].into_church() }
fn list3() -> Term { vec![1, 2, 3].into_church() }
fn   gt1() -> Term { app!(C(), gt(), 1.into_church()) }

#[bench]
fn church_null(b: &mut Bencher) {
    b.iter(|| { beta(app(null(), list3()), HAP, 0) } );
}

#[bench]
fn church_head(b: &mut Bencher) {
    b.iter(|| { beta(app(head(), list3()), HAP, 0) } );
}

#[bench]
fn church_tail(b: &mut Bencher) {
    b.iter(|| { beta(app(tail(), list3()), HAP, 0) } );
}

#[bench]
fn church_length(b: &mut Bencher) {
    b.iter(|| { beta(app(length(), list3()), HAP, 0) } );
}

#[bench]
fn church_list(b: &mut Bencher) {
    b.iter(|| { beta(app!(list(), 3.into_church(), 1.into_church(), 2.into_church(), 3.into_church()), HAP, 0) } );
}

#[bench]
fn church_reverse(b: &mut Bencher) {
    b.iter(|| { beta(app(reverse(), list3()), HAP, 0) } );
}

#[bench]
fn church_append(b: &mut Bencher) {
    b.iter(|| { beta(app!(append(), list1(), list2()), HAP, 0) } );
}

#[bench]
fn church_index(b: &mut Bencher) {
    b.iter(|| { beta(app!(index(), 1.into_church(), list3()), HAP, 0) } );
}

#[bench]
fn church_map(b: &mut Bencher) {
    b.iter(|| { beta(app!(map(), gt1(), list3()), HAP, 0) } );
}

#[bench]
fn church_foldl(b: &mut Bencher) {
    b.iter(|| { beta(app!(foldl(), add(), 0.into_church(), list3()), HAP, 0) } );
}

#[bench]
fn church_foldr(b: &mut Bencher) {
    b.iter(|| { beta(app!(foldr(), add(), 0.into_church(), list3()), HAP, 0) } );
}

#[bench]
fn church_filter(b: &mut Bencher) {
    b.iter(|| { beta(app!(filter(), gt1(), list3()), HAP, 0) } );
}

#[bench]
fn church_init(b: &mut Bencher) {
    b.iter(|| { beta(app(init(), list3()), HAP, 0) } );
}

#[bench]
fn church_last(b: &mut Bencher) {
    b.iter(|| { beta(app(last(), list3()), HAP, 0) } );
}

#[bench]
fn church_zip(b: &mut Bencher) {
    b.iter(|| { beta(app!(zip(), list3(), list3()), HAP, 0) } );
}

#[bench]
fn church_zip_with(b: &mut Bencher) {
    b.iter(|| { beta(app!(zip_with(), fls(), list3(), list3()), HAP, 0) } );
}

#[bench]
fn church_take(b: &mut Bencher) {
    b.iter(|| { beta(app!(take(), 2.into_church(), list3()), HAP, 0) } );
}

#[bench]
fn church_take_while(b: &mut Bencher) {
    b.iter(|| { beta(app!(take_while(), is_zero(), list3()), HAP, 0) } );
}
<commit_msg>fixed is_nil bench<commit_after>#![cfg(feature = "encoding")]

#![feature(test)]
extern crate test;

#[macro_use]
extern crate lambda_calculus as lambda;

use test::Bencher;
use lambda::church::list::*;
use lambda::church::numerals::*;
use lambda::church::boolean::fls;
use lambda::combinators::C;
use lambda::*;

fn list1() -> Term { vec![1].into_church() }
fn list2() -> Term { vec![2, 3].into_church() }
fn list3() -> Term { vec![1, 2, 3].into_church() }
fn   gt1() -> Term { app!(C(), gt(), 1.into_church()) }

#[bench]
fn church_is_nil(b: &mut Bencher) {
    b.iter(|| { beta(app(is_nil(), list3()), HAP, 0) } );
}

#[bench]
fn church_head(b: &mut Bencher) {
    b.iter(|| { beta(app(head(), list3()), HAP, 0) } );
}

#[bench]
fn church_tail(b: &mut Bencher) {
    b.iter(|| { beta(app(tail(), list3()), HAP, 0) } );
}

#[bench]
fn church_length(b: &mut Bencher) {
    b.iter(|| { beta(app(length(), list3()), HAP, 0) } );
}

#[bench]
fn church_list(b: &mut Bencher) {
    b.iter(|| { beta(app!(list(), 3.into_church(), 1.into_church(), 2.into_church(), 3.into_church()), HAP, 0) } );
}

#[bench]
fn church_reverse(b: &mut Bencher) {
    b.iter(|| { beta(app(reverse(), list3()), HAP, 0) } );
}

#[bench]
fn church_append(b: &mut Bencher) {
    b.iter(|| { beta(app!(append(), list1(), list2()), HAP, 0) } );
}

#[bench]
fn church_index(b: &mut Bencher) {
    b.iter(|| { beta(app!(index(), 1.into_church(), list3()), HAP, 0) } );
}

#[bench]
fn church_map(b: &mut Bencher) {
    b.iter(|| { beta(app!(map(), gt1(), list3()), HAP, 0) } );
}

#[bench]
fn church_foldl(b: &mut Bencher) {
    b.iter(|| { beta(app!(foldl(), add(), 0.into_church(), list3()), HAP, 0) } );
}

#[bench]
fn church_foldr(b: &mut Bencher) {
    b.iter(|| { beta(app!(foldr(), add(), 0.into_church(), list3()), HAP, 0) } );
}

#[bench]
fn church_filter(b: &mut Bencher) {
    b.iter(|| { beta(app!(filter(), gt1(), list3()), HAP, 0) } );
}

#[bench]
fn church_init(b: &mut Bencher) {
    b.iter(|| { beta(app(init(), list3()), HAP, 0) } );
}

#[bench]
fn church_last(b: &mut Bencher) {
    b.iter(|| { beta(app(last(), list3()), HAP, 0) } );
}

#[bench]
fn church_zip(b: &mut Bencher) {
    b.iter(|| { beta(app!(zip(), list3(), list3()), HAP, 0) } );
}

#[bench]
fn church_zip_with(b: &mut Bencher) {
    b.iter(|| { beta(app!(zip_with(), fls(), list3(), list3()), HAP, 0) } );
}

#[bench]
fn church_take(b: &mut Bencher) {
    b.iter(|| { beta(app!(take(), 2.into_church(), list3()), HAP, 0) } );
}

#[bench]
fn church_take_while(b: &mut Bencher) {
    b.iter(|| { beta(app!(take_while(), is_zero(), list3()), HAP, 0) } );
}
<|endoftext|>


<commit_before>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use stb_image::bindgen::*;

use std::any::Any;
use std::libc;
use std::libc::{c_void, c_int};
use std::vec::raw::from_buf_raw;

pub struct Image<T> {
    width   : uint,
    height  : uint,
    depth   : uint,
    data    : ~[T],
}

pub fn new_image<T>(width: uint, height: uint, depth: uint, data: ~[T]) -> Image<T> {
    Image::<T> {
        width   : width,
        height  : height,
        depth   : depth,
        data    : data,
    }
}

pub enum LoadResult {
    Error(~str),
    ImageU8(Image<u8>),
    ImageF32(Image<f32>),
}

impl LoadResult {
    pub fn from_result(res: Result<LoadResult,~Any>)-> LoadResult {
        match res {
            Ok(res) => res,
            Err(e)  => Error(e.to_str()),
        }
    }
}

pub fn load(path: ~str) -> LoadResult {
    let force_depth = 0;
    load_with_depth(path, force_depth, false)
}

fn load_internal<T>(buf : *T, w : c_int, h : c_int, d : c_int) -> Image<T> {
    unsafe {
        // FIXME: Shouldn't copy; instead we should use a sendable resource. They
        // aren't particularly safe yet though.
        let data = from_buf_raw(buf, (w * h * d) as uint);
        libc::free(buf as *mut c_void);
        Image::<T>{
            width   : w as uint,
            height  : h as uint,
            depth   : d as uint,
            data    : data}
    }
}

pub fn load_with_depth(path: ~str, force_depth: uint, convert_hdr:bool) -> LoadResult {
    unsafe {
        let mut width = 0 as c_int;
        let mut height = 0 as c_int;
        let mut depth = 0 as c_int;
        path.to_c_str().with_ref(|bytes| {
            if !convert_hdr && stbi_is_hdr(bytes)!=0   {
                let buffer = stbi_loadf(bytes,
                                        &mut width,
                                        &mut height,
                                        &mut depth,
                                        force_depth as c_int);
                if buffer.is_null() {
                    Error(~"stbi_loadf failed")
                } else {
                    ImageF32(load_internal(buffer, width, height, depth))
                }
            } else {
                let buffer = stbi_load(bytes,
                                       &mut width,
                                       &mut height,
                                       &mut depth,
                                       force_depth as c_int);
                if buffer.is_null() {
                    Error(~"stbi_load failed")
                } else {
                    ImageU8(load_internal(buffer, width, height, depth))
                }
            }
        })
    }
}

pub fn load_from_memory(buffer: &[u8]) -> LoadResult {
    let force_depth = 0;
    load_from_memory_with_depth(buffer, force_depth, false)
}

pub fn load_from_memory_with_depth(buffer: &[u8], force_depth: uint, convert_hdr:bool) -> LoadResult {
    unsafe {
        let mut width = 0 as c_int;
        let mut height = 0 as c_int;
        let mut depth = 0 as c_int;
        if !convert_hdr && stbi_is_hdr_from_memory(buffer.as_ptr(), buffer.len() as c_int) != 0 {
            let buffer = stbi_loadf_from_memory(buffer.as_ptr(),
                                                buffer.len() as c_int,
                                                &mut width,
                                                &mut height,
                                                &mut depth,
                                                force_depth as c_int);
            if buffer.is_null() {
                Error(~"stbi_loadf_from_memory failed")
            } else {
                let actual_depth = if force_depth != 0 { force_depth as c_int } else { depth };
                ImageF32(load_internal(buffer, width, height, actual_depth))
            }
        } else {
            let buffer = stbi_load_from_memory(buffer.as_ptr(),
                                               buffer.len() as c_int,
                                               &mut width,
                                               &mut height,
                                               &mut depth,
                                               force_depth as c_int);
            if buffer.is_null() {
                Error(~"stbi_load_from_memory failed")
            } else {
                let actual_depth = if force_depth != 0 { force_depth as c_int } else { depth };
                ImageU8(load_internal(buffer, width, height, actual_depth))
            }
        }
    }
}
<commit_msg>Updated to latest Rust: std::vec -> std::slice<commit_after>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use stb_image::bindgen::*;

use std::any::Any;
use std::libc;
use std::libc::{c_void, c_int};
use std::slice::raw::from_buf_raw;

pub struct Image<T> {
    width   : uint,
    height  : uint,
    depth   : uint,
    data    : ~[T],
}

pub fn new_image<T>(width: uint, height: uint, depth: uint, data: ~[T]) -> Image<T> {
    Image::<T> {
        width   : width,
        height  : height,
        depth   : depth,
        data    : data,
    }
}

pub enum LoadResult {
    Error(~str),
    ImageU8(Image<u8>),
    ImageF32(Image<f32>),
}

impl LoadResult {
    pub fn from_result(res: Result<LoadResult,~Any>)-> LoadResult {
        match res {
            Ok(res) => res,
            Err(e)  => Error(e.to_str()),
        }
    }
}

pub fn load(path: ~str) -> LoadResult {
    let force_depth = 0;
    load_with_depth(path, force_depth, false)
}

fn load_internal<T>(buf : *T, w : c_int, h : c_int, d : c_int) -> Image<T> {
    unsafe {
        // FIXME: Shouldn't copy; instead we should use a sendable resource. They
        // aren't particularly safe yet though.
        let data = from_buf_raw(buf, (w * h * d) as uint);
        libc::free(buf as *mut c_void);
        Image::<T>{
            width   : w as uint,
            height  : h as uint,
            depth   : d as uint,
            data    : data}
    }
}

pub fn load_with_depth(path: ~str, force_depth: uint, convert_hdr:bool) -> LoadResult {
    unsafe {
        let mut width = 0 as c_int;
        let mut height = 0 as c_int;
        let mut depth = 0 as c_int;
        path.to_c_str().with_ref(|bytes| {
            if !convert_hdr && stbi_is_hdr(bytes)!=0   {
                let buffer = stbi_loadf(bytes,
                                        &mut width,
                                        &mut height,
                                        &mut depth,
                                        force_depth as c_int);
                if buffer.is_null() {
                    Error(~"stbi_loadf failed")
                } else {
                    ImageF32(load_internal(buffer, width, height, depth))
                }
            } else {
                let buffer = stbi_load(bytes,
                                       &mut width,
                                       &mut height,
                                       &mut depth,
                                       force_depth as c_int);
                if buffer.is_null() {
                    Error(~"stbi_load failed")
                } else {
                    ImageU8(load_internal(buffer, width, height, depth))
                }
            }
        })
    }
}

pub fn load_from_memory(buffer: &[u8]) -> LoadResult {
    let force_depth = 0;
    load_from_memory_with_depth(buffer, force_depth, false)
}

pub fn load_from_memory_with_depth(buffer: &[u8], force_depth: uint, convert_hdr:bool) -> LoadResult {
    unsafe {
        let mut width = 0 as c_int;
        let mut height = 0 as c_int;
        let mut depth = 0 as c_int;
        if !convert_hdr && stbi_is_hdr_from_memory(buffer.as_ptr(), buffer.len() as c_int) != 0 {
            let buffer = stbi_loadf_from_memory(buffer.as_ptr(),
                                                buffer.len() as c_int,
                                                &mut width,
                                                &mut height,
                                                &mut depth,
                                                force_depth as c_int);
            if buffer.is_null() {
                Error(~"stbi_loadf_from_memory failed")
            } else {
                let actual_depth = if force_depth != 0 { force_depth as c_int } else { depth };
                ImageF32(load_internal(buffer, width, height, actual_depth))
            }
        } else {
            let buffer = stbi_load_from_memory(buffer.as_ptr(),
                                               buffer.len() as c_int,
                                               &mut width,
                                               &mut height,
                                               &mut depth,
                                               force_depth as c_int);
            if buffer.is_null() {
                Error(~"stbi_load_from_memory failed")
            } else {
                let actual_depth = if force_depth != 0 { force_depth as c_int } else { depth };
                ImageU8(load_internal(buffer, width, height, actual_depth))
            }
        }
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright 2019 Sanguk Park
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use core::mem::MaybeUninit;

use crate::addr::*;
use crate::arch::*;
use crate::boot_params::*;
use crate::cpu::*;
use crate::load::*;
use crate::memiter::*;
use crate::mm::*;
use crate::mpool::*;
use crate::page::*;
use crate::types::*;
use crate::vm::*;

extern "C" {
    fn plat_console_init();
    fn arch_one_time_init();
    fn dlog_enable_lock();

    /// The stack to be used by the CPUs.
    static callstacks: [[u8; STACK_SIZE]; MAX_CPUS];

    /// A record for boot CPU. Its field `stack_bottom` is initialized.
    /// Hafnium loader writes booted CPU ID on `cpus.id` and initializes the CPU
    /// stack by the address in `cpus.stack_bottom`.
    /// (See src/arch/aarch64/hypervisor/plat_entry.S and cpu_entry.S.)
    ///
    /// Initializing static variables with pointers in Rust failed here. We left
    /// the initialization code of `cpus` in `src/cpu.c`.
    static boot_cpu: Cpu;
}

/// TODO(HfO2): Refactor api.rs and remove `pub` (#44.)
pub struct Hypervisor {
    pub mpool: MPool,
    pub memory_manager: MemoryManager,
    pub cpu_manager: CpuManager,
    pub vm_manager: VmManager,
}

/// Note(HfO2): this variable was originally of type
/// MaybeUninit<[u8; mem::size_of::<RawPageTable>() * HEAP_PAGES]>,
/// but it was not aligned to PAGE_SIZE.
static mut PTABLE_BUF: MaybeUninit<[RawPage; HEAP_PAGES]> = MaybeUninit::uninit();

static mut INITED: bool = false;

/// A singleton collecting all managers in Hafnium.
///
/// This is dependency-free; Typical solutions of mutable and shared static
/// objects delay their initialization. Considering concurrency, they often use
/// std::sync features to prevent racy initialization. But Hafnium is different.
///
///  - The initialization is _always_ happened once in the specific time.
///  - During the time, no other thread is running; Hafnium runs as if it were
///    a single-thread program.
///  - After the initialization, Hafnium may make a non-exclusive reference of
///    singletons, but they have their own way for Hafnium to safely write them.
///
/// Therefore, I do not use a safe wrapper for initialization such as
/// `std::sync::Once` and `lazy_static`.
///
/// TODO(HfO2): This `pub` is required by mm_init, which is only used by unit
/// tests. Resolving #46 may help to remove the `pub`.
pub static mut HAFNIUM: MaybeUninit<Hypervisor> = MaybeUninit::uninit();

/// Performs one-time initialisation of the hypervisor.
#[no_mangle]
unsafe extern "C" fn one_time_init(c: *mut Cpu) -> *mut Cpu {
    if INITED {
        return c;
    }

    // Make sure the console is initialised before calling dlog.
    plat_console_init();

    dlog!("Initialising hafnium\n");

    arch_one_time_init();
    arch_cpu_module_init();

    let mut ppool = MPool::new();
    ppool.free_pages(Pages::from_raw(
        PTABLE_BUF.get_mut().as_mut_ptr(),
        HEAP_PAGES,
    ));

    let mm = MemoryManager::new(&ppool).expect("mm_init failed");

    // We have to initialize memory manager early, because
    // Stage2::invalidate_tlb refers global state. TODO(HfO2): Refactor the
    // function to recieve the info from caller. Maybe MemoryManager should
    // have a wrapper function for page table.
    HAFNIUM.get_mut().memory_manager = mm;
    let mm = &HAFNIUM.get_ref().memory_manager;

    // Enable locks now that mm is initialised.
    dlog_enable_lock();
    mpool_enable_locks();

    let mut hypervisor_ptable = mm.hypervisor_ptable.lock();

    let params = boot_params_get(&mut hypervisor_ptable, &mut ppool)
        .expect("unable to retrieve boot params");

    let cpum = CpuManager::new(
        &params.cpu_ids[..params.cpu_count],
        boot_cpu.id,
        &callstacks,
    );

    for i in 0..params.mem_ranges_count {
        dlog!(
            "Memory range: 0x{:x} - 0x{:x}\n",
            pa_addr(params.mem_ranges[i].begin),
            pa_addr(params.mem_ranges[i].end) - 1
        );
    }

    dlog!(
        "Ramdisk range: 0x{:x} - 0x{:x}\n",
        pa_addr(params.initrd_begin),
        pa_addr(params.initrd_end) - 1
    );

    // Map initrd in, and initialise cpio parser.
    hypervisor_ptable
        .identity_map(params.initrd_begin, params.initrd_end, Mode::R, &ppool)
        .expect("unable to map initrd in");

    let initrd = pa_addr(params.initrd_begin) as *mut _;
    let cpio = MemIter::from_raw(
        initrd,
        pa_difference(params.initrd_begin, params.initrd_end),
    );

    HAFNIUM.get_mut().vm_manager = VmManager::new();

    // Load all VMs.
    let primary_initrd = load_primary(
        &mut HAFNIUM.get_mut().vm_manager,
        &mut hypervisor_ptable,
        &cpio,
        params.kernel_arg,
        &ppool,
    )
    .expect("unable to load primary VM");

    // load_secondary will add regions assigned to the secondary VMs from
    // mem_ranges to reserved_ranges.
    let mut update: BootParamsUpdate = BootParamsUpdate::new(
        pa_from_va(va_from_ptr(primary_initrd.next as usize as *const _)),
        pa_from_va(va_from_ptr(primary_initrd.limit as usize as *const _)),
    );

    load_secondary(
        &mut HAFNIUM.get_mut().vm_manager,
        &mut hypervisor_ptable,
        &cpio,
        &params,
        &mut update,
        &mut ppool,
    )
    .expect("unable to load secondary VMs");

    // Prepare to run by updating bootparams as seen by primary VM.
    boot_params_update(&mut hypervisor_ptable, &mut update, &mut ppool)
        .expect("plat_update_boot_params failed");

    hypervisor_ptable.defrag(&ppool);

    // Initialise HAFNIUM.
    HAFNIUM.get_mut().mpool = ppool;
    HAFNIUM.get_mut().cpu_manager = cpum;

    // Enable TLB invalidation for VM page table updates.
    mm_vm_enable_invalidation();

    dlog!("Hafnium initialisation completed\n");
    INITED = true;

    hafnium().cpu_manager.boot_cpu()

    // From now on, other pCPUs are on in order to run multiple vCPUs. Thus
    // you may safely make readonly references to the Hafnium singleton, but
    // may not modify the singleton without proper locking.
}

pub fn hafnium() -> &'static Hypervisor {
    unsafe { HAFNIUM.get_ref() }
}

// The entry point of CPUs when they are turned on. It is supposed to initialise
// all state and return the first vCPU to run.
#[no_mangle]
pub unsafe extern "C" fn cpu_main(c: *const Cpu) -> *mut VCpu {
    if !mm_cpu_init() {
        panic!("mm_cpu_init failed");
    }

    let primary = hafnium().vm_manager.get(HF_PRIMARY_VM_ID).unwrap();
    let vcpu = primary.vcpus.get(cpu_index(&*c)).unwrap();
    let vm = vcpu.vm;

    // TODO(HfO2): vcpu needs to be borrowed exclusively, which is safe but
    // discouraged. Move this code into one_time_init().
    let vcpu = vcpu as *const _ as usize as *mut VCpu;
    (*vcpu).set_cpu(c);

    // Reset the registers to give a clean start for the primary's vCPU.
    (*vcpu).inner.get_mut().regs.reset(true, &*vm, (*c).id);

    vcpu
}
<commit_msg>Use ptr::write more.<commit_after>/*
 * Copyright 2019 Sanguk Park
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use core::mem::MaybeUninit;
use core::ptr;

use crate::addr::*;
use crate::arch::*;
use crate::boot_params::*;
use crate::cpu::*;
use crate::load::*;
use crate::memiter::*;
use crate::mm::*;
use crate::mpool::*;
use crate::page::*;
use crate::types::*;
use crate::vm::*;

extern "C" {
    fn plat_console_init();
    fn arch_one_time_init();
    fn dlog_enable_lock();

    /// The stack to be used by the CPUs.
    static callstacks: [[u8; STACK_SIZE]; MAX_CPUS];

    /// A record for boot CPU. Its field `stack_bottom` is initialized.
    /// Hafnium loader writes booted CPU ID on `cpus.id` and initializes the CPU
    /// stack by the address in `cpus.stack_bottom`.
    /// (See src/arch/aarch64/hypervisor/plat_entry.S and cpu_entry.S.)
    ///
    /// Initializing static variables with pointers in Rust failed here. We left
    /// the initialization code of `cpus` in `src/cpu.c`.
    static boot_cpu: Cpu;
}

/// TODO(HfO2): Refactor api.rs and remove `pub` (#44.)
pub struct Hypervisor {
    pub mpool: MPool,
    pub memory_manager: MemoryManager,
    pub cpu_manager: CpuManager,
    pub vm_manager: VmManager,
}

/// Note(HfO2): this variable was originally of type
/// MaybeUninit<[u8; mem::size_of::<RawPageTable>() * HEAP_PAGES]>,
/// but it was not aligned to PAGE_SIZE.
static mut PTABLE_BUF: MaybeUninit<[RawPage; HEAP_PAGES]> = MaybeUninit::uninit();

static mut INITED: bool = false;

/// A singleton collecting all managers in Hafnium.
///
/// This is dependency-free; Typical solutions of mutable and shared static
/// objects delay their initialization. Considering concurrency, they often use
/// std::sync features to prevent racy initialization. But Hafnium is different.
///
///  - The initialization is _always_ happened once in the specific time.
///  - During the time, no other thread is running; Hafnium runs as if it were
///    a single-thread program.
///  - After the initialization, Hafnium may make a non-exclusive reference of
///    singletons, but they have their own way for Hafnium to safely write them.
///
/// Therefore, I do not use a safe wrapper for initialization such as
/// `std::sync::Once` and `lazy_static`.
///
/// TODO(HfO2): This `pub` is required by mm_init, which is only used by unit
/// tests. Resolving #46 may help to remove the `pub`.
pub static mut HAFNIUM: MaybeUninit<Hypervisor> = MaybeUninit::uninit();

/// Performs one-time initialisation of the hypervisor.
#[no_mangle]
unsafe extern "C" fn one_time_init(c: *mut Cpu) -> *mut Cpu {
    if INITED {
        return c;
    }

    // Make sure the console is initialised before calling dlog.
    plat_console_init();

    dlog!("Initialising hafnium\n");

    arch_one_time_init();
    arch_cpu_module_init();

    let mut ppool = MPool::new();
    ppool.free_pages(Pages::from_raw(
        PTABLE_BUF.get_mut().as_mut_ptr(),
        HEAP_PAGES,
    ));

    let mm = MemoryManager::new(&ppool).expect("mm_init failed");

    // We have to initialize memory manager early, because
    // Stage2::invalidate_tlb refers global state. TODO(HfO2): Refactor the
    // function to recieve the info from caller. Maybe MemoryManager should
    // have a wrapper function for page table.
    ptr::write(&mut HAFNIUM.get_mut().memory_manager, mm);
    let mm = &HAFNIUM.get_ref().memory_manager;

    // Enable locks now that mm is initialised.
    dlog_enable_lock();
    mpool_enable_locks();

    let mut hypervisor_ptable = mm.hypervisor_ptable.lock();

    let params = boot_params_get(&mut hypervisor_ptable, &mut ppool)
        .expect("unable to retrieve boot params");

    let cpum = CpuManager::new(
        &params.cpu_ids[..params.cpu_count],
        boot_cpu.id,
        &callstacks,
    );

    for i in 0..params.mem_ranges_count {
        dlog!(
            "Memory range: 0x{:x} - 0x{:x}\n",
            pa_addr(params.mem_ranges[i].begin),
            pa_addr(params.mem_ranges[i].end) - 1
        );
    }

    dlog!(
        "Ramdisk range: 0x{:x} - 0x{:x}\n",
        pa_addr(params.initrd_begin),
        pa_addr(params.initrd_end) - 1
    );

    // Map initrd in, and initialise cpio parser.
    hypervisor_ptable
        .identity_map(params.initrd_begin, params.initrd_end, Mode::R, &ppool)
        .expect("unable to map initrd in");

    let initrd = pa_addr(params.initrd_begin) as *mut _;
    let cpio = MemIter::from_raw(
        initrd,
        pa_difference(params.initrd_begin, params.initrd_end),
    );

    ptr::write(&mut HAFNIUM.get_mut().vm_manager, VmManager::new());

    // Load all VMs.
    let primary_initrd = load_primary(
        &mut HAFNIUM.get_mut().vm_manager,
        &mut hypervisor_ptable,
        &cpio,
        params.kernel_arg,
        &ppool,
    )
    .expect("unable to load primary VM");

    // load_secondary will add regions assigned to the secondary VMs from
    // mem_ranges to reserved_ranges.
    let mut update: BootParamsUpdate = BootParamsUpdate::new(
        pa_from_va(va_from_ptr(primary_initrd.next as usize as *const _)),
        pa_from_va(va_from_ptr(primary_initrd.limit as usize as *const _)),
    );

    load_secondary(
        &mut HAFNIUM.get_mut().vm_manager,
        &mut hypervisor_ptable,
        &cpio,
        &params,
        &mut update,
        &mut ppool,
    )
    .expect("unable to load secondary VMs");

    // Prepare to run by updating bootparams as seen by primary VM.
    boot_params_update(&mut hypervisor_ptable, &mut update, &mut ppool)
        .expect("plat_update_boot_params failed");

    hypervisor_ptable.defrag(&ppool);

    // Initialise HAFNIUM.
    ptr::write(&mut HAFNIUM.get_mut().mpool, ppool);
    ptr::write(&mut HAFNIUM.get_mut().cpu_manager, cpum);

    // Enable TLB invalidation for VM page table updates.
    mm_vm_enable_invalidation();

    dlog!("Hafnium initialisation completed\n");
    INITED = true;

    hafnium().cpu_manager.boot_cpu()

    // From now on, other pCPUs are on in order to run multiple vCPUs. Thus
    // you may safely make readonly references to the Hafnium singleton, but
    // may not modify the singleton without proper locking.
}

pub fn hafnium() -> &'static Hypervisor {
    unsafe { HAFNIUM.get_ref() }
}

// The entry point of CPUs when they are turned on. It is supposed to initialise
// all state and return the first vCPU to run.
#[no_mangle]
pub unsafe extern "C" fn cpu_main(c: *const Cpu) -> *mut VCpu {
    if !mm_cpu_init() {
        panic!("mm_cpu_init failed");
    }

    let primary = hafnium().vm_manager.get(HF_PRIMARY_VM_ID).unwrap();
    let vcpu = primary.vcpus.get(cpu_index(&*c)).unwrap();
    let vm = vcpu.vm;

    // TODO(HfO2): vcpu needs to be borrowed exclusively, which is safe but
    // discouraged. Move this code into one_time_init().
    let vcpu = vcpu as *const _ as usize as *mut VCpu;
    (*vcpu).set_cpu(c);

    // Reset the registers to give a clean start for the primary's vCPU.
    (*vcpu).inner.get_mut().regs.reset(true, &*vm, (*c).id);

    vcpu
}
<|endoftext|>


<commit_before>//! Implements `ThreadExecutor`.

use std::time::Duration;
use std::cell::RefCell;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver, TryRecvError};
use uci::{SetOption, OptionDescription};
use search::*;

/// Turns a `SearchThread` into `SearchExecutor`.
pub struct ThreadExecutor<T: Search> {
    tt: Arc<T::HashTable>,
    messages_tx: Sender<String>,
    reports_rx: Receiver<SearchReport<T::ReportData>>,
    reports_tx: Sender<SearchReport<T::ReportData>>,
    pending_report: RefCell<Option<SearchReport<T::ReportData>>>,
}

impl<T: Search> SearchExecutor for ThreadExecutor<T> {
    type HashTable = T::HashTable;

    type SearchNode = T::SearchNode;

    type ReportData = T::ReportData;

    fn new(tt: Arc<Self::HashTable>) -> Self {
        let (messages_tx, _) = channel();
        let (reports_tx, reports_rx) = channel();
        Self {
            tt: tt,
            messages_tx: messages_tx,
            reports_rx: reports_rx,
            reports_tx: reports_tx,
            pending_report: RefCell::new(None),
        }
    }

    fn start_search(&mut self, params: SearchParams<Self::SearchNode>) {
        let (messages_tx, messages_rx) = channel();
        self.messages_tx = messages_tx;
        T::start_thread(params,
                        self.tt.clone(),
                        self.reports_tx.clone(),
                        messages_rx);
    }

    fn wait_report(&self, timeout_after: Duration) {
        let mut report = self.pending_report.borrow_mut();
        if report.is_none() {
            *report = self.reports_rx.recv_timeout(timeout_after).ok();
        }
    }

    fn try_recv_report(&mut self) -> Result<SearchReport<Self::ReportData>, TryRecvError> {
        self.pending_report
            .borrow_mut()
            .take()
            .ok_or(TryRecvError::Empty)
            .or_else(|_| self.reports_rx.try_recv())
    }

    fn send_message(&mut self, msg: &str) {
        self.messages_tx.send(msg.to_string()).ok();
    }
}

impl<T: Search> SetOption for ThreadExecutor<T> {
    fn options() -> Vec<(String, OptionDescription)> {
        T::options()
    }

    fn set_option(name: &str, value: &str) {
        T::set_option(name, value);
    }
}
<commit_msg>minor code improvement<commit_after>//! Implements `ThreadExecutor`.

use std::time::Duration;
use std::cell::RefCell;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver, TryRecvError};
use uci::{SetOption, OptionDescription};
use search::*;

/// Turns a `SearchThread` into `SearchExecutor`.
pub struct ThreadExecutor<T: Search> {
    tt: Arc<T::HashTable>,
    messages_tx: Sender<String>,
    reports_rx: Receiver<SearchReport<T::ReportData>>,
    reports_tx: Sender<SearchReport<T::ReportData>>,
    pending_report: RefCell<Option<SearchReport<T::ReportData>>>,
}

impl<T: Search> SearchExecutor for ThreadExecutor<T> {
    type HashTable = T::HashTable;

    type SearchNode = T::SearchNode;

    type ReportData = T::ReportData;

    fn new(tt: Arc<Self::HashTable>) -> Self {
        let (reports_tx, reports_rx) = channel();
        Self {
            tt: tt,
            messages_tx: channel().0,
            reports_rx: reports_rx,
            reports_tx: reports_tx,
            pending_report: RefCell::new(None),
        }
    }

    fn start_search(&mut self, params: SearchParams<Self::SearchNode>) {
        let (messages_tx, messages_rx) = channel();
        self.messages_tx = messages_tx;
        T::start_thread(params,
                        self.tt.clone(),
                        self.reports_tx.clone(),
                        messages_rx);
    }

    fn wait_report(&self, timeout_after: Duration) {
        let mut report = self.pending_report.borrow_mut();
        if report.is_none() {
            *report = self.reports_rx.recv_timeout(timeout_after).ok();
        }
    }

    fn try_recv_report(&mut self) -> Result<SearchReport<Self::ReportData>, TryRecvError> {
        self.pending_report
            .borrow_mut()
            .take()
            .ok_or(TryRecvError::Empty)
            .or_else(|_| self.reports_rx.try_recv())
    }

    fn send_message(&mut self, msg: &str) {
        self.messages_tx.send(msg.to_string()).ok();
    }
}

impl<T: Search> SetOption for ThreadExecutor<T> {
    fn options() -> Vec<(String, OptionDescription)> {
        T::options()
    }

    fn set_option(name: &str, value: &str) {
        T::set_option(name, value);
    }
}
<|endoftext|>


<commit_before>use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::collections::HashMap;
use std::collections::HashSet;
use std::io::BufRead;
use std::io::BufReader;
use std::io::Read;

use common::manhattan_distance;
use common::Solution;

type Coordinate = (usize, usize);

#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Copy, Clone)]
struct Unit {
    pos: Coordinate,
    hp: u8,
    power: u8,
    faction: char,
}

impl Default for Unit {
    fn default() -> Self {
        Unit {
            pos: (0, 0),
            hp: 200,
            power: 3,
            faction: 'E',
        }
    }
}

impl Unit {
    pub fn is_alive(&self) -> bool {
        self.hp > 0
    }
}

#[derive(Default)]
pub struct Day15 {
    walls: Vec<Vec<bool>>,
    units: Vec<Unit>,
    alive: [usize; 2],
}

impl Day15 {
    pub fn new() -> Self {
        Default::default()
    }

    fn read_input(&mut self, input: &mut Read) {
        let reader = BufReader::new(input);

        for (y, line) in reader.lines().enumerate() {
            let line = line.unwrap();
            let mut current = vec![false; line.len()];

            for (x, c) in line.chars().enumerate() {
                match c {
                    '#' => { current[x] = true; }
                    'E' | 'G' => {
                        self.units.push(Unit {
                            pos: (y, x),
                            faction: c,
                            ..Default::default()
                        });
                        self.alive[if c == 'E' { 0 } else { 1 }] += 1;
                    }
                    '.' => {}
                    c => panic!("Invalid tile {}!", c),
                }
            }
            self.walls.push(current);
        }
    }

    fn get_movement(&self, unit: usize) -> Option<Coordinate> {
        let initial = self.units[unit].pos;
        let faction = self.units[unit].faction;

        let enemy_positions: HashSet<Coordinate> = self.units.iter()
            .filter(|x| x.faction != faction && x.is_alive())
            .map(|x| x.pos).collect();

        let all_positions: HashSet<Coordinate> = self.units.iter()
            .filter(|x| x.is_alive())
            .map(|x| x.pos).collect();

        let mut todo = BinaryHeap::new();
        let mut prev: HashMap<Coordinate, Coordinate> = HashMap::new();

        todo.push(Reverse((0, initial)));

        while let Some(Reverse((d, (y, x)))) = todo.pop() {
            let next = [
                (y - 1, x),
                (y, x - 1),
                (y, x + 1),
                (y + 1, x),
            ];

            for pos in &next {
                if !prev.contains_key(pos) && !self.walls[pos.0][pos.1] {
                    if enemy_positions.contains(pos) {
                        return prev.remove(&(y, x));
                    } else if !all_positions.contains(pos) {
                        ;
                        let prev_step = *prev.get(&(y, x)).unwrap_or(pos);
                        prev.insert(*pos, prev_step);
                        todo.push(Reverse((d + 1, *pos)));
                    }
                }
            }
        }
        None
    }

    fn get_attack(&self, unit: usize) -> Option<usize> {
        let initial = self.units[unit].pos;
        let faction = self.units[unit].faction;

        let to_attack = self.units.iter()
            .enumerate()
            .filter(|(_, x)| x.faction != faction && x.is_alive())
            .filter(|(_, x)| manhattan_distance(x.pos, initial) == 1)
            .min_by(|&(_, a), &(_, b)| a.hp.cmp(&b.hp).then(a.pos.cmp(&b.pos)));

        if let Some((index, _)) = to_attack {
            Some(index)
        } else {
            None
        }
    }

    fn simulate(&mut self) -> bool {
        self.units.sort_unstable();
        for i in 0..self.units.len() {
            if !self.units[i].is_alive() {
                continue;
            }

            if self.alive[0] == 0 || self.alive[1] == 0 {
                return false;
            }

            if let Some(new_pos) = self.get_movement(i) {
                self.units[i].pos = new_pos;
            }

            if let Some(target) = self.get_attack(i) {
                let power = self.units[i].power;
                let target = &mut self.units[target];
                target.hp = target.hp.saturating_sub(power);

                if target.hp == 0 {
                    match target.faction {
                        'E' => { self.alive[0] -= 1 }
                        'G' => { self.alive[1] -= 1 }
                        _ => panic!(),
                    };
                }
            }
        }

        true
    }

    #[allow(dead_code)]
    fn print(&self) {
        let positions: HashMap<_, _> = self.units.iter()
            .filter(|x| x.is_alive())
            .map(|x| (x.pos, x))
            .collect();

        for y in 0..self.walls.len() {
            let mut buf = String::new();
            let mut unit_buf = String::new();

            for x in 0..self.walls[y].len() {
                if let Some(unit) = positions.get(&(y, x)) {
                    buf.push(unit.faction);

                    unit_buf += &format!("    {}({})", unit.faction, unit.hp);
                } else if self.walls[y][x] {
                    buf.push('#');
                } else {
                    buf.push('.');
                }
            }
            println!("{}{}", buf, unit_buf);
        }
    }

    fn return_score(&self, rounds: usize) -> String {
        let result: usize = rounds * self.units.iter().map(|x| x.hp as usize)
            .sum::<usize>();
        format!("{}", result)
    }

    fn set_elf_power(&mut self, power: u8) {
        for unit in self.units.iter_mut() {
            if unit.faction == 'E' {
                unit.power = power;
            }
        }
    }
}

impl Solution for Day15 {
    fn part1(&mut self, input: &mut Read) -> String {
        self.read_input(input);
        let mut rounds = 0;
        while self.simulate() {
            rounds += 1;
        }
        self.return_score(rounds)
    }

    fn part2(&mut self, input: &mut Read) -> String {
        self.read_input(input);
        let backup = self.units.clone();
        let starting_elves = self.alive[0];
        let starting_goblins = self.alive[1];

        for power in 4..=255 {
            self.units = backup.clone();
            self.alive[0] = starting_elves;
            self.alive[1] = starting_goblins;

            self.set_elf_power(power);
            let mut rounds = 0;
            while self.simulate() {
                rounds += 1;
            }

            if self.alive[0] == starting_elves {
                return self.return_score(rounds);
            }

        }
        unreachable!()
    }
}

#[cfg(test)]
mod tests {
    use common::Solution;
    use day15::Day15;

    const SAMPLE_INPUT: [&[u8]; 6] = [
        include_bytes!("samples/15.1.txt"),
        include_bytes!("samples/15.2.txt"),
        include_bytes!("samples/15.3.txt"),
        include_bytes!("samples/15.4.txt"),
        include_bytes!("samples/15.5.txt"),
        include_bytes!("samples/15.6.txt"),
    ];

    const SAMPLE1_OUTPUT: [&str; 6] = [
        "27730",
        "36334",
        "39514",
        "27755",
        "28944",
        "18740",
    ];

    #[test]
    fn sample_part1() {
        for (input, output) in SAMPLE_INPUT.iter().zip(SAMPLE1_OUTPUT.iter()) {
            let mut instance = Day15::new();
            assert_eq!(*output, instance.part1(&mut input.as_ref()));
        }
    }

    #[test]
    fn sample_part2() {
        let indices = [0, 2, 3, 4, 5];
        let outputs = [
            "4988",
            "31284",
            "3478",
            "6474",
            "1140",
        ];

        for (&input, output) in indices.iter().zip(outputs.iter()) {
            let mut instance = Day15::new();
            assert_eq!(*output, instance.part2(&mut SAMPLE_INPUT[input]));
        }
    }
}
<commit_msg>Small improvement to day 15.<commit_after>use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::collections::HashMap;
use std::collections::HashSet;
use std::io::BufRead;
use std::io::BufReader;
use std::io::Read;

use common::manhattan_distance;
use common::Solution;

type Coordinate = (usize, usize);

#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Copy, Clone)]
struct Unit {
    pos: Coordinate,
    hp: u8,
    power: u8,
    faction: char,
}

impl Default for Unit {
    fn default() -> Self {
        Unit {
            pos: (0, 0),
            hp: 200,
            power: 3,
            faction: 'E',
        }
    }
}

impl Unit {
    pub fn is_alive(&self) -> bool {
        self.hp > 0
    }
}

#[derive(Default)]
pub struct Day15 {
    walls: Vec<Vec<bool>>,
    units: Vec<Unit>,
    alive: [usize; 2],
}

impl Day15 {
    pub fn new() -> Self {
        Default::default()
    }

    fn read_input(&mut self, input: &mut Read) {
        let reader = BufReader::new(input);

        for (y, line) in reader.lines().enumerate() {
            let line = line.unwrap();
            let mut current = vec![false; line.len()];

            for (x, c) in line.chars().enumerate() {
                match c {
                    '#' => { current[x] = true; }
                    'E' | 'G' => {
                        self.units.push(Unit {
                            pos: (y, x),
                            faction: c,
                            ..Default::default()
                        });
                        self.alive[if c == 'E' { 0 } else { 1 }] += 1;
                    }
                    '.' => {}
                    c => panic!("Invalid tile {}!", c),
                }
            }
            self.walls.push(current);
        }
    }

    fn get_movement(&self, unit: usize) -> Option<Coordinate> {
        let initial = self.units[unit].pos;
        let faction = self.units[unit].faction;

        let enemy_positions: HashSet<Coordinate> = self.units.iter()
            .filter(|x| x.faction != faction && x.is_alive())
            .map(|x| x.pos).collect();

        let all_positions: HashSet<Coordinate> = self.units.iter()
            .filter(|x| x.is_alive())
            .map(|x| x.pos).collect();

        let mut todo = BinaryHeap::new();
        let mut prev: HashMap<Coordinate, Coordinate> = HashMap::new();

        todo.push(Reverse((0, initial)));

        while let Some(Reverse((d, (y, x)))) = todo.pop() {
            let next = [
                (y - 1, x),
                (y, x - 1),
                (y, x + 1),
                (y + 1, x),
            ];

            for pos in &next {
                if !prev.contains_key(pos) && !self.walls[pos.0][pos.1] {
                    if enemy_positions.contains(pos) {
                        return prev.remove(&(y, x));
                    } else if !all_positions.contains(pos) {
                        ;
                        let prev_step = *prev.get(&(y, x)).unwrap_or(pos);
                        prev.insert(*pos, prev_step);
                        todo.push(Reverse((d + 1, *pos)));
                    }
                }
            }
        }
        None
    }

    fn get_attack(&self, unit: usize) -> Option<usize> {
        let initial = self.units[unit].pos;
        let faction = self.units[unit].faction;

        let to_attack = self.units.iter()
            .enumerate()
            .filter(|(_, x)| x.faction != faction && x.is_alive())
            .filter(|(_, x)| manhattan_distance(x.pos, initial) == 1)
            .min_by(|&(_, a), &(_, b)| a.hp.cmp(&b.hp).then(a.pos.cmp(&b.pos)));

        if let Some((index, _)) = to_attack {
            Some(index)
        } else {
            None
        }
    }

    fn simulate(&mut self) -> bool {
        self.units.sort_unstable();
        for i in 0..self.units.len() {
            if !self.units[i].is_alive() {
                continue;
            }

            if self.alive[0] == 0 || self.alive[1] == 0 {
                return false;
            }

            if let Some(new_pos) = self.get_movement(i) {
                self.units[i].pos = new_pos;
            }

            if let Some(target) = self.get_attack(i) {
                let power = self.units[i].power;
                let target = &mut self.units[target];
                target.hp = target.hp.saturating_sub(power);

                if target.hp == 0 {
                    match target.faction {
                        'E' => { self.alive[0] -= 1 }
                        'G' => { self.alive[1] -= 1 }
                        _ => panic!(),
                    };
                }
            }
        }

        true
    }

    #[allow(dead_code)]
    fn print(&self) {
        let positions: HashMap<_, _> = self.units.iter()
            .filter(|x| x.is_alive())
            .map(|x| (x.pos, x))
            .collect();

        for y in 0..self.walls.len() {
            let mut buf = String::new();
            let mut unit_buf = String::new();

            for x in 0..self.walls[y].len() {
                if let Some(unit) = positions.get(&(y, x)) {
                    buf.push(unit.faction);

                    unit_buf += &format!("    {}({})", unit.faction, unit.hp);
                } else if self.walls[y][x] {
                    buf.push('#');
                } else {
                    buf.push('.');
                }
            }
            println!("{}{}", buf, unit_buf);
        }
    }

    fn return_score(&self, rounds: usize) -> String {
        let result: usize = rounds * self.units.iter().map(|x| x.hp as usize)
            .sum::<usize>();
        format!("{}", result)
    }

    fn set_elf_power(&mut self, power: u8) {
        for unit in self.units.iter_mut() {
            if unit.faction == 'E' {
                unit.power = power;
            }
        }
    }
}

impl Solution for Day15 {
    fn part1(&mut self, input: &mut Read) -> String {
        self.read_input(input);
        let mut rounds = 0;
        while self.simulate() {
            rounds += 1;
        }
        self.return_score(rounds)
    }

    fn part2(&mut self, input: &mut Read) -> String {
        self.read_input(input);
        let backup = self.units.clone();
        let starting_elves = self.alive[0];
        let starting_goblins = self.alive[1];

        let mut power = 4;

        loop {
            self.units = backup.clone();
            self.alive[0] = starting_elves;
            self.alive[1] = starting_goblins;

            self.set_elf_power(power);
            let mut rounds = 0;
            while self.simulate() {
                rounds += 1;
            }

            if self.alive[0] == starting_elves {
                return self.return_score(rounds);
            }

            let to_kill = 200 / power;
            while to_kill == 200 / power {
                power += 1;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use common::Solution;
    use day15::Day15;

    const SAMPLE_INPUT: [&[u8]; 6] = [
        include_bytes!("samples/15.1.txt"),
        include_bytes!("samples/15.2.txt"),
        include_bytes!("samples/15.3.txt"),
        include_bytes!("samples/15.4.txt"),
        include_bytes!("samples/15.5.txt"),
        include_bytes!("samples/15.6.txt"),
    ];

    const SAMPLE1_OUTPUT: [&str; 6] = [
        "27730",
        "36334",
        "39514",
        "27755",
        "28944",
        "18740",
    ];

    #[test]
    fn sample_part1() {
        for (input, output) in SAMPLE_INPUT.iter().zip(SAMPLE1_OUTPUT.iter()) {
            let mut instance = Day15::new();
            assert_eq!(*output, instance.part1(&mut input.as_ref()));
        }
    }

    #[test]
    fn sample_part2() {
        let indices = [0, 2, 3, 4, 5];
        let outputs = [
            "4988",
            "31284",
            "3478",
            "6474",
            "1140",
        ];

        for (&input, output) in indices.iter().zip(outputs.iter()) {
            let mut instance = Day15::new();
            assert_eq!(*output, instance.part2(&mut SAMPLE_INPUT[input]));
        }
    }
}
<|endoftext|>


<commit_before>#![deny(warnings)]

extern crate cargo_registry;
extern crate postgres;
extern crate semver;
extern crate time;

use std::env;
use std::collections::HashMap;
use time::Duration;

use cargo_registry::{VersionDownload, Version, Model};

static LIMIT: i64 = 1000;

#[allow(dead_code)] // dead in tests
fn main() {
    let daemon = env::args().nth(1).as_ref().map(|s| &s[..])
                    == Some("daemon");
    let sleep = env::args().nth(2).map(|s| s.parse::<u32>().unwrap());
    loop {
        let conn = postgres::Connection::connect(&env("DATABASE_URL")[..],
                                                 &postgres::SslMode::None).unwrap();
        update(&conn).unwrap();
        drop(conn);
        if daemon {
            std::thread::sleep_ms(sleep.unwrap() * 1000);
        } else {
            break
        }
    }
}

fn env(s: &str) -> String {
    match env::var(s).ok() {
        Some(s) => s,
        None => panic!("must have `{}` defined", s),
    }
}

fn update(conn: &postgres::GenericConnection) -> postgres::Result<()> {
    let mut max = 0;
    loop {
        // FIXME(rust-lang/rust#27401): weird declaration to make sure this
        // variable gets dropped.
        let tx; tx = try!(conn.transaction());
        {
            let stmt = try!(tx.prepare("SELECT * FROM version_downloads \
                                        WHERE processed = FALSE AND id > $1
                                        ORDER BY id ASC
                                        LIMIT $2"));
            let mut rows = try!(stmt.query(&[&max, &LIMIT]));
            match try!(collect(&tx, &mut rows)) {
                None => break,
                Some(m) => max = m,
            }
        }
        tx.set_commit();
        try!(tx.finish());
    }
    Ok(())
}

fn collect(tx: &postgres::Transaction,
           rows: &mut postgres::rows::Rows) -> postgres::Result<Option<i32>> {

    // Anything older than 24 hours ago will be frozen and will not be queried
    // against again.
    let cutoff = time::now_utc().to_timespec();
    let cutoff = cutoff + Duration::days(-1);

    let mut map = HashMap::new();
    for row in rows.iter() {
        let download: VersionDownload = Model::from_row(&row);
        assert!(map.insert(download.id, download).is_none());
    }
    println!("updating {} versions (cutoff {})", map.len(),
             time::at(cutoff).rfc822());
    if map.len() == 0 { return Ok(None) }

    let mut max = 0;
    let mut total = 0;
    for (id, download) in map.iter() {
        if *id > max { max = *id; }
        if download.date > cutoff && download.counted == download.downloads {
            continue
        }
        let amt = download.downloads - download.counted;

        let crate_id = Version::find(tx, download.version_id).unwrap().crate_id;

        // Update the total number of version downloads
        try!(tx.execute("UPDATE versions
                         SET downloads = downloads + $1
                         WHERE id = $2",
                        &[&amt, &download.version_id]));
        // Update the total number of crate downloads
        try!(tx.execute("UPDATE crates SET downloads = downloads + $1
                         WHERE id = $2", &[&amt, &crate_id]));

        // Update the total number of crate downloads for today
        let cnt = try!(tx.execute("UPDATE crate_downloads
                                   SET downloads = downloads + $2
                                   WHERE crate_id = $1 AND date = date($3)",
                                  &[&crate_id, &amt, &download.date]));
        if cnt == 0 {
            try!(tx.execute("INSERT INTO crate_downloads
                             (crate_id, downloads, date)
                             VALUES ($1, $2, $3)",
                            &[&crate_id, &amt, &download.date]));
        }

        // Flag this row as having been processed if we're passed the cutoff,
        // and unconditionally increment the number of counted downloads.
        try!(tx.execute("UPDATE version_downloads
                         SET processed = $2, counted = counted + $3
                         WHERE id = $1",
                        &[id, &(download.date < cutoff), &amt]));
        total += amt as i64;
    }

    // After everything else is done, update the global counter of total
    // downloads.
    try!(tx.execute("UPDATE metadata SET total_downloads = total_downloads + $1",
                    &[&total]));

    Ok(Some(max))
}

#[cfg(test)]
mod test {
    use std::collections::HashMap;

    use postgres;
    use semver;

    use cargo_registry::{Version, Crate, User};

    fn conn() -> postgres::Connection {
        postgres::Connection::connect(&::env("TEST_DATABASE_URL")[..],
                                      &postgres::SslMode::None).unwrap()
    }

    fn user(conn: &postgres::Transaction) -> User{
        User::find_or_insert(conn, "login", None, None, None,
                             "access_token", "api_token").unwrap()
    }

    fn crate_downloads(tx: &postgres::Transaction, id: i32, expected: usize) {
        let stmt = tx.prepare("SELECT * FROM crate_downloads
                               WHERE crate_id = $1").unwrap();
        let dl: i32 = stmt.query(&[&id]).unwrap().iter()
                          .next().unwrap().get("downloads");
        assert_eq!(dl, expected as i32);
    }

    #[test]
    fn increment() {
        let conn = conn();
        let tx = conn.transaction().unwrap();
        let user = user(&tx);
        let krate = Crate::find_or_insert(&tx, "foo", user.id, &None, &None,
                                          &None, &None, &[], &None, &None,
                                          &None).unwrap();
        let version = Version::insert(&tx, krate.id,
                                      &semver::Version::parse("1.0.0").unwrap(),
                                      &HashMap::new(), &[]).unwrap();
        tx.execute("INSERT INTO version_downloads \
                    (version_id, downloads, counted, date, processed)
                    VALUES ($1, 1, 0, current_date, false)",
                   &[&version.id]).unwrap();
        tx.execute("INSERT INTO version_downloads \
                    (version_id, downloads, counted, date, processed)
                    VALUES ($1, 1, 0, current_date, true)",
                   &[&version.id]).unwrap();
        ::update(&tx).unwrap();
        assert_eq!(Version::find(&tx, version.id).unwrap().downloads, 1);
        assert_eq!(Crate::find(&tx, krate.id).unwrap().downloads, 1);
        crate_downloads(&tx, krate.id, 1);
        ::update(&tx).unwrap();
        assert_eq!(Version::find(&tx, version.id).unwrap().downloads, 1);
    }

    #[test]
    fn set_processed_true() {
        let conn = conn();
        let tx = conn.transaction().unwrap();
        let user = user(&tx);
        let krate = Crate::find_or_insert(&tx, "foo", user.id, &None,
                                          &None, &None, &None, &[], &None,
                                          &None, &None).unwrap();
        let version = Version::insert(&tx, krate.id,
                                      &semver::Version::parse("1.0.0").unwrap(),
                                      &HashMap::new(), &[]).unwrap();
        tx.execute("INSERT INTO version_downloads \
                    (version_id, downloads, counted, date, processed)
                    VALUES ($1, 2, 2, current_date - interval '2 days', false)",
                   &[&version.id]).unwrap();
        ::update(&tx).unwrap();
        let stmt = tx.prepare("SELECT processed FROM version_downloads
                               WHERE version_id = $1").unwrap();
        let processed: bool = stmt.query(&[&version.id]).unwrap().iter()
                                  .next().unwrap().get("processed");
        assert!(processed);
    }

    #[test]
    fn increment_a_little() {
        let conn = conn();
        let tx = conn.transaction().unwrap();
        let user = user(&tx);
        let krate = Crate::find_or_insert(&tx, "foo", user.id, &None,
                                          &None, &None, &None, &[], &None,
                                          &None, &None).unwrap();
        let version = Version::insert(&tx, krate.id,
                                      &semver::Version::parse("1.0.0").unwrap(),
                                      &HashMap::new(), &[]).unwrap();
        tx.execute("INSERT INTO version_downloads \
                    (version_id, downloads, counted, date, processed)
                    VALUES ($1, 2, 1, current_date, false)",
                   &[&version.id]).unwrap();
        tx.execute("INSERT INTO version_downloads \
                    (version_id, downloads, counted, date, processed)
                    VALUES ($1, 1, 0, current_date, false)",
                   &[&version.id]).unwrap();
        ::update(&tx).unwrap();
        assert_eq!(Version::find(&tx, version.id).unwrap().downloads, 2);
        assert_eq!(Crate::find(&tx, krate.id).unwrap().downloads, 2);
        crate_downloads(&tx, krate.id, 2);
        ::update(&tx).unwrap();
        assert_eq!(Version::find(&tx, version.id).unwrap().downloads, 2);
    }
}
<commit_msg>Don't use sleep_ms<commit_after>#![deny(warnings)]

extern crate cargo_registry;
extern crate postgres;
extern crate semver;
extern crate time;

use std::env;
use std::collections::HashMap;
use std::time::Duration;

use cargo_registry::{VersionDownload, Version, Model};

static LIMIT: i64 = 1000;

#[allow(dead_code)] // dead in tests
fn main() {
    let daemon = env::args().nth(1).as_ref().map(|s| &s[..])
                    == Some("daemon");
    let sleep = env::args().nth(2).map(|s| s.parse().unwrap());
    loop {
        let conn = postgres::Connection::connect(&env("DATABASE_URL")[..],
                                                 &postgres::SslMode::None).unwrap();
        update(&conn).unwrap();
        drop(conn);
        if daemon {
            std::thread::sleep(Duration::new(sleep.unwrap(), 0));
        } else {
            break
        }
    }
}

fn env(s: &str) -> String {
    match env::var(s).ok() {
        Some(s) => s,
        None => panic!("must have `{}` defined", s),
    }
}

fn update(conn: &postgres::GenericConnection) -> postgres::Result<()> {
    let mut max = 0;
    loop {
        // FIXME(rust-lang/rust#27401): weird declaration to make sure this
        // variable gets dropped.
        let tx; tx = try!(conn.transaction());
        {
            let stmt = try!(tx.prepare("SELECT * FROM version_downloads \
                                        WHERE processed = FALSE AND id > $1
                                        ORDER BY id ASC
                                        LIMIT $2"));
            let mut rows = try!(stmt.query(&[&max, &LIMIT]));
            match try!(collect(&tx, &mut rows)) {
                None => break,
                Some(m) => max = m,
            }
        }
        tx.set_commit();
        try!(tx.finish());
    }
    Ok(())
}

fn collect(tx: &postgres::Transaction,
           rows: &mut postgres::rows::Rows) -> postgres::Result<Option<i32>> {
    use time::Duration;

    // Anything older than 24 hours ago will be frozen and will not be queried
    // against again.
    let cutoff = time::now_utc().to_timespec();
    let cutoff = cutoff + Duration::days(1);

    let mut map = HashMap::new();
    for row in rows.iter() {
        let download: VersionDownload = Model::from_row(&row);
        assert!(map.insert(download.id, download).is_none());
    }
    println!("updating {} versions (cutoff {})", map.len(),
             time::at(cutoff).rfc822());
    if map.len() == 0 { return Ok(None) }

    let mut max = 0;
    let mut total = 0;
    for (id, download) in map.iter() {
        if *id > max { max = *id; }
        if download.date > cutoff && download.counted == download.downloads {
            continue
        }
        let amt = download.downloads - download.counted;

        let crate_id = Version::find(tx, download.version_id).unwrap().crate_id;

        // Update the total number of version downloads
        try!(tx.execute("UPDATE versions
                         SET downloads = downloads + $1
                         WHERE id = $2",
                        &[&amt, &download.version_id]));
        // Update the total number of crate downloads
        try!(tx.execute("UPDATE crates SET downloads = downloads + $1
                         WHERE id = $2", &[&amt, &crate_id]));

        // Update the total number of crate downloads for today
        let cnt = try!(tx.execute("UPDATE crate_downloads
                                   SET downloads = downloads + $2
                                   WHERE crate_id = $1 AND date = date($3)",
                                  &[&crate_id, &amt, &download.date]));
        if cnt == 0 {
            try!(tx.execute("INSERT INTO crate_downloads
                             (crate_id, downloads, date)
                             VALUES ($1, $2, $3)",
                            &[&crate_id, &amt, &download.date]));
        }

        // Flag this row as having been processed if we're passed the cutoff,
        // and unconditionally increment the number of counted downloads.
        try!(tx.execute("UPDATE version_downloads
                         SET processed = $2, counted = counted + $3
                         WHERE id = $1",
                        &[id, &(download.date < cutoff), &amt]));
        total += amt as i64;
    }

    // After everything else is done, update the global counter of total
    // downloads.
    try!(tx.execute("UPDATE metadata SET total_downloads = total_downloads + $1",
                    &[&total]));

    Ok(Some(max))
}

#[cfg(test)]
mod test {
    use std::collections::HashMap;

    use postgres;
    use semver;

    use cargo_registry::{Version, Crate, User};

    fn conn() -> postgres::Connection {
        postgres::Connection::connect(&::env("TEST_DATABASE_URL")[..],
                                      &postgres::SslMode::None).unwrap()
    }

    fn user(conn: &postgres::Transaction) -> User{
        User::find_or_insert(conn, "login", None, None, None,
                             "access_token", "api_token").unwrap()
    }

    fn crate_downloads(tx: &postgres::Transaction, id: i32, expected: usize) {
        let stmt = tx.prepare("SELECT * FROM crate_downloads
                               WHERE crate_id = $1").unwrap();
        let dl: i32 = stmt.query(&[&id]).unwrap().iter()
                          .next().unwrap().get("downloads");
        assert_eq!(dl, expected as i32);
    }

    #[test]
    fn increment() {
        let conn = conn();
        let tx = conn.transaction().unwrap();
        let user = user(&tx);
        let krate = Crate::find_or_insert(&tx, "foo", user.id, &None, &None,
                                          &None, &None, &[], &None, &None,
                                          &None).unwrap();
        let version = Version::insert(&tx, krate.id,
                                      &semver::Version::parse("1.0.0").unwrap(),
                                      &HashMap::new(), &[]).unwrap();
        tx.execute("INSERT INTO version_downloads \
                    (version_id, downloads, counted, date, processed)
                    VALUES ($1, 1, 0, current_date, false)",
                   &[&version.id]).unwrap();
        tx.execute("INSERT INTO version_downloads \
                    (version_id, downloads, counted, date, processed)
                    VALUES ($1, 1, 0, current_date, true)",
                   &[&version.id]).unwrap();
        ::update(&tx).unwrap();
        assert_eq!(Version::find(&tx, version.id).unwrap().downloads, 1);
        assert_eq!(Crate::find(&tx, krate.id).unwrap().downloads, 1);
        crate_downloads(&tx, krate.id, 1);
        ::update(&tx).unwrap();
        assert_eq!(Version::find(&tx, version.id).unwrap().downloads, 1);
    }

    #[test]
    fn set_processed_true() {
        let conn = conn();
        let tx = conn.transaction().unwrap();
        let user = user(&tx);
        let krate = Crate::find_or_insert(&tx, "foo", user.id, &None,
                                          &None, &None, &None, &[], &None,
                                          &None, &None).unwrap();
        let version = Version::insert(&tx, krate.id,
                                      &semver::Version::parse("1.0.0").unwrap(),
                                      &HashMap::new(), &[]).unwrap();
        tx.execute("INSERT INTO version_downloads \
                    (version_id, downloads, counted, date, processed)
                    VALUES ($1, 2, 2, current_date - interval '2 days', false)",
                   &[&version.id]).unwrap();
        ::update(&tx).unwrap();
        let stmt = tx.prepare("SELECT processed FROM version_downloads
                               WHERE version_id = $1").unwrap();
        let processed: bool = stmt.query(&[&version.id]).unwrap().iter()
                                  .next().unwrap().get("processed");
        assert!(processed);
    }

    #[test]
    fn increment_a_little() {
        let conn = conn();
        let tx = conn.transaction().unwrap();
        let user = user(&tx);
        let krate = Crate::find_or_insert(&tx, "foo", user.id, &None,
                                          &None, &None, &None, &[], &None,
                                          &None, &None).unwrap();
        let version = Version::insert(&tx, krate.id,
                                      &semver::Version::parse("1.0.0").unwrap(),
                                      &HashMap::new(), &[]).unwrap();
        tx.execute("INSERT INTO version_downloads \
                    (version_id, downloads, counted, date, processed)
                    VALUES ($1, 2, 1, current_date, false)",
                   &[&version.id]).unwrap();
        tx.execute("INSERT INTO version_downloads \
                    (version_id, downloads, counted, date, processed)
                    VALUES ($1, 1, 0, current_date, false)",
                   &[&version.id]).unwrap();
        ::update(&tx).unwrap();
        assert_eq!(Version::find(&tx, version.id).unwrap().downloads, 2);
        assert_eq!(Crate::find(&tx, krate.id).unwrap().downloads, 2);
        crate_downloads(&tx, krate.id, 2);
        ::update(&tx).unwrap();
        assert_eq!(Version::find(&tx, version.id).unwrap().downloads, 2);
    }
}
<|endoftext|>