<commit_before>
use super::effective_address;
use super::Core;

pub fn ay_pd_8(core: &mut Core) -> u32 {
	let ea = effective_address::predecrement_ay(core);
	core.read_data_byte(ea)
}
pub fn ax_pd_8(core: &mut Core) -> (u32, u32) {
	let ea = effective_address::predecrement_ax(core);
	(core.read_data_byte(ea), ea)
}
pub fn ay_pi_8(core: &mut Core) -> u32 {
	let ea = effective_address::postincrement_ay(core);
	core.read_data_byte(ea)
}
pub fn ay_ai_8(core: &mut Core) -> u32 {
	let ea = effective_address::address_indirect_ay(core);
	core.read_data_byte(ea)
}
pub fn ay_di_8(core: &mut Core) -> u32 {
	let ea = effective_address::displacement_ay(core);
	core.read_data_byte(ea)
}
pub fn ay_ix_8(core: &mut Core) -> u32 {
	let ea = effective_address::index_ay(core);
	core.read_data_byte(ea)
}
pub fn aw_8(core: &mut Core) -> u32 {
	let ea = effective_address::absolute_word(core);
	core.read_data_byte(ea)
}
pub fn al_8(core: &mut Core) -> u32 {
	let ea = effective_address::absolute_long(core);
	core.read_data_byte(ea)
}
pub fn pcdi_8(core: &mut Core) -> u32 {
	let ea = effective_address::displacement_pc(core);
	core.read_program_byte(ea)
}
pub fn pcix_8(core: &mut Core) -> u32 {
	let ea = effective_address::index_pc(core);
	core.read_program_byte(ea)
}
pub fn imm_8(core: &mut Core) -> u32 {
	let extension = core.read_imm_u16();
	mask_out_above_8!(extension) as u32
}
pub fn dx(core: &mut Core) -> u32 {
	dx!(core)
}
pub fn dy(core: &mut Core) -> u32 {
	dy!(core)
}
pub fn ax(core: &mut Core) -> u32 {
	ax!(core)
}
pub fn ay(core: &mut Core) -> u32 {
	ay!(core)
}
#[cfg(test)]
mod tests {
	use super::super::Core;
	use super::{ax_pd_8, ay_pd_8};
	use ram::{AddressBus, SUPERVISOR_DATA};

	#[test]
	fn predecrement_ax() {
		let mut core = Core::new(0x40);
		for i in 0..8 {
			let addr: u32 = 0x200 + 4*i;
			core.dar[8+i as usize] = addr;
			// write just before where A0-A7 points
			let adjustment = if i == 7 {2} else {1};
			core.mem.write_byte(SUPERVISOR_DATA, addr - adjustment, 0x11*i);
		}
		core.ir = 0b1111_1001_1111_1010; // X=4, Y=2
		let core = &mut core;

		assert_eq!(512+4*4, core.dar[8+4]);
		let (dst, ea) = ax_pd_8(core);
		assert_eq!(0x44, dst);
		assert_eq!(512+4*4-1, core.dar[8+4]);
		assert_eq!(512+4*4-1, ea);

		core.ir = 0b1111_1111_1111_1111; // X=7, Y=7
		assert_eq!(512+4*7, core.dar[8+7]);
		let (dst, ea) = ax_pd_8(core);
		assert_eq!(0x77, dst);
		// A7 is kept even
		assert_eq!(512+4*7-2, core.dar[8+7]);
		assert_eq!(512+4*7-2, ea);
	}
	#[test]
	fn predecrement_ay() {
		let mut core = Core::new(0x40);
		for i in 0..8 {
			let addr: u32 = 0x200 + 4*i;
			core.dar[8+i as usize] = addr;
			// write just before where A0-A7 points
			let adjustment = if i == 7 {2} else {1};
			core.mem.write_byte(SUPERVISOR_DATA, addr - adjustment, 0x11*i);
		}

		core.ir = 0b1111_1001_1111_1010; // X=4, Y=2
		let core = &mut core;
		assert_eq!(512+4*2, core.dar[8+2]);
		assert_eq!(0x22, ay_pd_8(core));
		assert_eq!(512+4*2-1, core.dar[8+2]);

		core.ir = 0b1111_1011_1111_1111; // X=5, Y=7
		assert_eq!(512+4*7, core.dar[8+7]);
		assert_eq!(0x77, ay_pd_8(core));
		// A7 is kept even
		assert_eq!(512+4*7-2, core.dar[8+7]);
	}
}
<commit_msg>Add 16-bit operators<commit_after>
use super::effective_address;
use super::Core;

pub fn ay_pd_8(core: &mut Core) -> u32 {
	let ea = effective_address::predecrement_ay(core);
	core.read_data_byte(ea)
}
pub fn ax_pd_8(core: &mut Core) -> (u32, u32) {
	let ea = effective_address::predecrement_ax(core);
	(core.read_data_byte(ea), ea)
}
pub fn ay_pi_8(core: &mut Core) -> u32 {
	let ea = effective_address::postincrement_ay(core);
	core.read_data_byte(ea)
}
pub fn ay_ai_8(core: &mut Core) -> u32 {
	let ea = effective_address::address_indirect_ay(core);
	core.read_data_byte(ea)
}
pub fn ay_di_8(core: &mut Core) -> u32 {
	let ea = effective_address::displacement_ay(core);
	core.read_data_byte(ea)
}
pub fn ay_ix_8(core: &mut Core) -> u32 {
	let ea = effective_address::index_ay(core);
	core.read_data_byte(ea)
}
pub fn aw_8(core: &mut Core) -> u32 {
	let ea = effective_address::absolute_word(core);
	core.read_data_byte(ea)
}
pub fn al_8(core: &mut Core) -> u32 {
	let ea = effective_address::absolute_long(core);
	core.read_data_byte(ea)
}
pub fn pcdi_8(core: &mut Core) -> u32 {
	let ea = effective_address::displacement_pc(core);
	core.read_program_byte(ea)
}
pub fn pcix_8(core: &mut Core) -> u32 {
	let ea = effective_address::index_pc(core);
	core.read_program_byte(ea)
}
pub fn imm_8(core: &mut Core) -> u32 {
	let extension = core.read_imm_u16();
	mask_out_above_8!(extension) as u32
}
pub fn ay_pd_16(core: &mut Core) -> u32 {
	let ea = effective_address::predecrement_ay(core);
	core.read_data_word(ea)
}
pub fn ax_pd_16(core: &mut Core) -> (u32, u32) {
	let ea = effective_address::predecrement_ax(core);
	(core.read_data_word(ea), ea)
}
pub fn ay_pi_16(core: &mut Core) -> u32 {
	let ea = effective_address::postincrement_ay(core);
	core.read_data_word(ea)
}
pub fn ay_ai_16(core: &mut Core) -> u32 {
	let ea = effective_address::address_indirect_ay(core);
	core.read_data_word(ea)
}
pub fn ay_di_16(core: &mut Core) -> u32 {
	let ea = effective_address::displacement_ay(core);
	core.read_data_word(ea)
}
pub fn ay_ix_16(core: &mut Core) -> u32 {
	let ea = effective_address::index_ay(core);
	core.read_data_word(ea)
}
pub fn aw_16(core: &mut Core) -> u32 {
	let ea = effective_address::absolute_word(core);
	core.read_data_word(ea)
}
pub fn al_16(core: &mut Core) -> u32 {
	let ea = effective_address::absolute_long(core);
	core.read_data_word(ea)
}
pub fn pcdi_16(core: &mut Core) -> u32 {
	let ea = effective_address::displacement_pc(core);
	core.read_program_word(ea)
}
pub fn pcix_16(core: &mut Core) -> u32 {
	let ea = effective_address::index_pc(core);
	core.read_program_word(ea)
}
pub fn imm_16(core: &mut Core) -> u32 {
	let extension = core.read_imm_u16();
	mask_out_above_8!(extension) as u32
}
pub fn dx(core: &mut Core) -> u32 {
	dx!(core)
}
pub fn dy(core: &mut Core) -> u32 {
	dy!(core)
}
pub fn ax(core: &mut Core) -> u32 {
	ax!(core)
}
pub fn ay(core: &mut Core) -> u32 {
	ay!(core)
}
#[cfg(test)]
mod tests {
	use super::super::Core;
	use super::{ax_pd_8, ay_pd_8};
	use ram::{AddressBus, SUPERVISOR_DATA};

	#[test]
	fn predecrement_ax() {
		let mut core = Core::new(0x40);
		for i in 0..8 {
			let addr: u32 = 0x200 + 4*i;
			core.dar[8+i as usize] = addr;
			// write just before where A0-A7 points
			let adjustment = if i == 7 {2} else {1};
			core.mem.write_byte(SUPERVISOR_DATA, addr - adjustment, 0x11*i);
		}
		core.ir = 0b1111_1001_1111_1010; // X=4, Y=2
		let core = &mut core;

		assert_eq!(512+4*4, core.dar[8+4]);
		let (dst, ea) = ax_pd_8(core);
		assert_eq!(0x44, dst);
		assert_eq!(512+4*4-1, core.dar[8+4]);
		assert_eq!(512+4*4-1, ea);

		core.ir = 0b1111_1111_1111_1111; // X=7, Y=7
		assert_eq!(512+4*7, core.dar[8+7]);
		let (dst, ea) = ax_pd_8(core);
		assert_eq!(0x77, dst);
		// A7 is kept even
		assert_eq!(512+4*7-2, core.dar[8+7]);
		assert_eq!(512+4*7-2, ea);
	}
	#[test]
	fn predecrement_ay() {
		let mut core = Core::new(0x40);
		for i in 0..8 {
			let addr: u32 = 0x200 + 4*i;
			core.dar[8+i as usize] = addr;
			// write just before where A0-A7 points
			let adjustment = if i == 7 {2} else {1};
			core.mem.write_byte(SUPERVISOR_DATA, addr - adjustment, 0x11*i);
		}

		core.ir = 0b1111_1001_1111_1010; // X=4, Y=2
		let core = &mut core;
		assert_eq!(512+4*2, core.dar[8+2]);
		assert_eq!(0x22, ay_pd_8(core));
		assert_eq!(512+4*2-1, core.dar[8+2]);

		core.ir = 0b1111_1011_1111_1111; // X=5, Y=7
		assert_eq!(512+4*7, core.dar[8+7]);
		assert_eq!(0x77, ay_pd_8(core));
		// A7 is kept even
		assert_eq!(512+4*7-2, core.dar[8+7]);
	}
}
<|endoftext|>


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

// Manages the SQL database that holds the list of registered domain names.
// Each records is made of the name, the private token and the Let's Encrypt
// challenge value.

use std::sync::mpsc::{channel, Receiver};
use std::thread;

use r2d2_sqlite::SqliteConnectionManager;
use r2d2;
use rusqlite::Row;
use rusqlite::Result as SqlResult;
use rusqlite::types::{ToSql, ToSqlOutput};

#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct DomainRecord {
    pub name: String,
    pub token: String,
    pub dns_challenge: Option<String>,
    pub local_ip: Option<String>,
    pub public_ip: Option<String>,
    pub timestamp: i64,
}

macro_rules! sqlstr {
    ($row:ident, $index:expr) => (
        {
            let raw = $row.get::<i32, String>($index);
            if raw.is_empty() {
                None
            } else {
                Some(raw)
            }
        }
    )
}

impl DomainRecord {
    fn from_sql(row: Row) -> Self {
        DomainRecord {
            name: row.get(0),
            token: row.get(1),
            dns_challenge: sqlstr!(row, 2),
            local_ip: sqlstr!(row, 3),
            public_ip: sqlstr!(row, 4),
            timestamp: row.get(5),
        }
    }

    pub fn new(name: &str,
               token: &str,
               dns_challenge: Option<&str>,
               local_ip: Option<&str>,
               public_ip: Option<&str>,
               timestamp: i64)
               -> Self {
        macro_rules! str2sql {
            ($val:expr) => (
                if $val.is_some() {
                    Some($val.unwrap().to_owned())
                } else {
                    None
                }
            )
        }

        DomainRecord {
            name: name.to_owned(),
            token: token.to_owned(),
            dns_challenge: str2sql!(dns_challenge),
            local_ip: str2sql!(local_ip),
            public_ip: str2sql!(public_ip),
            timestamp: timestamp,
        }
    }
}

unsafe impl Send for DomainRecord {}
unsafe impl Sync for DomainRecord {}

#[derive(Debug)]
pub enum DomainError {
    DbUnavailable,
    SQLError(String),
    NoRecord,
}

#[derive(Clone)]
pub struct DomainDb {
    pool: r2d2::Pool<SqliteConnectionManager>,
}

// try! like macro that sends back the error of a SQLite operation
// over the mpsc channel if needed.
macro_rules! sqltry {
    ($sql:expr, $tx:ident, $err:expr) => (
        match $sql {
            Err(_) => {
                $tx.send(Err($err)).unwrap();
                return;
            }
            Ok(value) => value,
        }
    );

    ($sql:expr, $tx:ident) => (
        match $sql {
            Err(err) => {
                $tx.send(Err(DomainError::SQLError(format!("{}", err)))).unwrap();
                return;
            }
            Ok(value) => value,
        }
    );
}

#[derive(Clone)]
pub enum SqlParam {
    Text(String),
    Integer(i64),
}

impl ToSql for SqlParam {
    fn to_sql(&self) -> SqlResult<ToSqlOutput> {
        match *self {
            SqlParam::Text(ref text) => text.to_sql(),
            SqlParam::Integer(ref number) => number.to_sql(),
        }
    }
}

impl DomainDb {
    pub fn new(path: &str) -> Self {
        let config = r2d2::Config::default();
        let manager = SqliteConnectionManager::new(path);
        let pool = r2d2::Pool::new(config, manager).expect("Unable to open database!");

        // Create the database table if needed.
        let conn = pool.get().unwrap();
        conn.execute("CREATE TABLE IF NOT EXISTS domains (
                      name          TEXT NOT NULL PRIMARY KEY,
                      token         TEXT NOT NULL,
                      dns_challenge TEXT NOT NULL,
                      local_ip      TEXT NOT NULL,
                      public_ip     TEXT NOT NULL,
                      timestamp     INTEGER)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains database: {}", err);
                            });

        conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS domains_token ON domains(token)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains.token index: {}", err);
                            });

        conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS domains_timestamp ON domains(timestamp)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains.timestamp index: {}", err);
                            });

        DomainDb { pool: pool }
    }

    fn select_record(&self,
                     request: &str,
                     value: &str)
                     -> Receiver<Result<DomainRecord, DomainError>> {
        let (tx, rx) = channel();

        // Run the sql command on a pooled thread.
        let pool = self.pool.clone();
        let value = value.to_owned();
        let request = request.to_owned();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            let mut stmt = sqltry!(conn.prepare(&request), tx);
            let mut rows = sqltry!(stmt.query(&[&value]), tx);
            if let Some(result_row) = rows.next() {
                let row = sqltry!(result_row, tx);
                tx.send(Ok(DomainRecord::from_sql(row))).unwrap();
            } else {
                tx.send(Err(DomainError::NoRecord)).unwrap();
            }
        });

        rx
    }

    pub fn get_record_by_name(&self, name: &str) -> Receiver<Result<DomainRecord, DomainError>> {
        self.select_record("SELECT name, token, dns_challenge, local_ip, public_ip, timestamp \
                            FROM domains WHERE name=$1",
                           name)
    }

    pub fn get_record_by_token(&self, token: &str) -> Receiver<Result<DomainRecord, DomainError>> {
        self.select_record("SELECT name, token, dns_challenge, local_ip, public_ip, timestamp \
                            FROM domains WHERE token=$1",
                           token)
    }

    pub fn add_record(&self, record: DomainRecord) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let record = record.clone();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            sqltry!(conn.execute("INSERT INTO domains VALUES ($1, $2, $3, $4, $5, $6)",
                                 &[&record.name,
                                   &record.token,
                                   &record.dns_challenge.unwrap_or("".to_owned()),
                                   &record.local_ip.unwrap_or("".to_owned()),
                                   &record.public_ip.unwrap_or("".to_owned()),
                                   &record.timestamp]),
                    tx);
            tx.send(Ok(())).unwrap();
        });

        rx
    }

    pub fn update_record(&self, record: DomainRecord) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let record = record.clone();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);

            sqltry!(conn.execute("UPDATE domains SET dns_challenge=$1, local_ip=$2, public_ip=$3, timestamp=$4 \
                                  WHERE name=$5 AND token=$6",
                                 &[&record.dns_challenge.unwrap_or("".to_owned()),
                                   &record.local_ip.unwrap_or("".to_owned()),
                                   &record.public_ip.unwrap_or("".to_owned()),
                                   &record.timestamp,
                                   &record.name,
                                   &record.token]),
                    tx);
            tx.send(Ok(())).unwrap();
        });

        rx
    }

    // Returns the number of rows affected.
    pub fn execute_1param_sql(&self,
                              request: &str,
                              value: SqlParam)
                              -> Receiver<Result<i32, DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let value = value.to_owned();
        let request = request.to_owned();
        thread::spawn(move || {
                          let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
                          let res = sqltry!(conn.execute(&request, &[&value]), tx);
                          tx.send(Ok(res)).unwrap();
                      });

        rx
    }

    pub fn delete_record_by_name(&self, name: &str) -> Receiver<Result<i32, DomainError>> {
        self.execute_1param_sql("DELETE FROM domains WHERE name=$1",
                                SqlParam::Text(name.to_owned()))
    }

    pub fn delete_record_by_token(&self, token: &str) -> Receiver<Result<i32, DomainError>> {
        self.execute_1param_sql("DELETE FROM domains WHERE token=$1",
                                SqlParam::Text(token.to_owned()))
    }

    pub fn flush(&self) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        thread::spawn(move || {
                          let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
                          sqltry!(conn.execute("DELETE FROM domains", &[]), tx);
                          tx.send(Ok(())).unwrap();
                      });
        rx
    }
}

#[test]
fn test_domain_store() {
    let db = DomainDb::new("domain_db_test.sqlite");

    // Start with an empty db.
    db.flush().recv().unwrap().expect("Flushing the db");

    // Check that we don't find any record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        Err(err) => panic!("Should not find a record by name in an empty db. {:?}", err),
        _ => panic!("Should not find a record by name in an empty db."),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find a record by token in an empty db."),
    }

    // Add a record without a dns challenge.
    let no_challenge_record =
        DomainRecord::new("test.example.org", "test-token", None, None, None, 0);
    db.add_record(no_challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Adding the no_challenge record");

    // Check that we can find it and that it matches our record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Ok(record) => assert_eq!(record, no_challenge_record),
        Err(err) => panic!("Failed to find record by name: {:?}", err),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Ok(record) => assert_eq!(record, no_challenge_record),
        Err(err) => panic!("Failed to find record by token: {:?}", err),
    }

    // Update the record to have challenge.
    let challenge_record = DomainRecord::new("test.example.org",
                                             "test-token",
                                             Some("dns-challenge"),
                                             None,
                                             None,
                                             0);
    db.update_record(challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Updating the challenge record");

    // Check that we can find it and that it matches our record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Ok(record) => assert_eq!(record, challenge_record),
        Err(err) => panic!("Failed to find record by name: {:?}", err),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Ok(record) => assert_eq!(record, challenge_record),
        Err(err) => panic!("Failed to find record by token: {:?}", err),
    }

    // Remove by name.
    db.delete_record_by_name(&challenge_record.name)
        .recv()
        .unwrap()
        .expect("Should delete record by name");
    match db.get_record_by_name(&challenge_record.name)
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find this record anymore."),
    }

    // Add again, and remove by token.
    db.add_record(challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Adding the challenge record");
    db.delete_record_by_token(&challenge_record.token)
        .recv()
        .unwrap()
        .expect("Should delete record by token");
    match db.get_record_by_name(&challenge_record.name)
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find this record anymore."),
    }
}
<commit_msg>Don't serialize None fields<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Manages the SQL database that holds the list of registered domain names.
// Each records is made of the name, the private token and the Let's Encrypt
// challenge value.

use std::sync::mpsc::{channel, Receiver};
use std::thread;

use r2d2_sqlite::SqliteConnectionManager;
use r2d2;
use rusqlite::Row;
use rusqlite::Result as SqlResult;
use rusqlite::types::{ToSql, ToSqlOutput};

#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct DomainRecord {
    pub name: String,
    pub token: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub dns_challenge: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub local_ip: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub public_ip: Option<String>,
    pub timestamp: i64,
}

macro_rules! sqlstr {
    ($row:ident, $index:expr) => (
        {
            let raw = $row.get::<i32, String>($index);
            if raw.is_empty() {
                None
            } else {
                Some(raw)
            }
        }
    )
}

impl DomainRecord {
    fn from_sql(row: Row) -> Self {
        DomainRecord {
            name: row.get(0),
            token: row.get(1),
            dns_challenge: sqlstr!(row, 2),
            local_ip: sqlstr!(row, 3),
            public_ip: sqlstr!(row, 4),
            timestamp: row.get(5),
        }
    }

    pub fn new(name: &str,
               token: &str,
               dns_challenge: Option<&str>,
               local_ip: Option<&str>,
               public_ip: Option<&str>,
               timestamp: i64)
               -> Self {
        macro_rules! str2sql {
            ($val:expr) => (
                if $val.is_some() {
                    Some($val.unwrap().to_owned())
                } else {
                    None
                }
            )
        }

        DomainRecord {
            name: name.to_owned(),
            token: token.to_owned(),
            dns_challenge: str2sql!(dns_challenge),
            local_ip: str2sql!(local_ip),
            public_ip: str2sql!(public_ip),
            timestamp: timestamp,
        }
    }
}

unsafe impl Send for DomainRecord {}
unsafe impl Sync for DomainRecord {}

#[derive(Debug)]
pub enum DomainError {
    DbUnavailable,
    SQLError(String),
    NoRecord,
}

#[derive(Clone)]
pub struct DomainDb {
    pool: r2d2::Pool<SqliteConnectionManager>,
}

// try! like macro that sends back the error of a SQLite operation
// over the mpsc channel if needed.
macro_rules! sqltry {
    ($sql:expr, $tx:ident, $err:expr) => (
        match $sql {
            Err(_) => {
                $tx.send(Err($err)).unwrap();
                return;
            }
            Ok(value) => value,
        }
    );

    ($sql:expr, $tx:ident) => (
        match $sql {
            Err(err) => {
                $tx.send(Err(DomainError::SQLError(format!("{}", err)))).unwrap();
                return;
            }
            Ok(value) => value,
        }
    );
}

#[derive(Clone)]
pub enum SqlParam {
    Text(String),
    Integer(i64),
}

impl ToSql for SqlParam {
    fn to_sql(&self) -> SqlResult<ToSqlOutput> {
        match *self {
            SqlParam::Text(ref text) => text.to_sql(),
            SqlParam::Integer(ref number) => number.to_sql(),
        }
    }
}

impl DomainDb {
    pub fn new(path: &str) -> Self {
        let config = r2d2::Config::default();
        let manager = SqliteConnectionManager::new(path);
        let pool = r2d2::Pool::new(config, manager).expect("Unable to open database!");

        // Create the database table if needed.
        let conn = pool.get().unwrap();
        conn.execute("CREATE TABLE IF NOT EXISTS domains (
                      name          TEXT NOT NULL PRIMARY KEY,
                      token         TEXT NOT NULL,
                      dns_challenge TEXT NOT NULL,
                      local_ip      TEXT NOT NULL,
                      public_ip     TEXT NOT NULL,
                      timestamp     INTEGER)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains database: {}", err);
                            });

        conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS domains_token ON domains(token)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains.token index: {}", err);
                            });

        conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS domains_timestamp ON domains(timestamp)",
                     &[])
            .unwrap_or_else(|err| {
                                panic!("Unable to create the domains.timestamp index: {}", err);
                            });

        DomainDb { pool: pool }
    }

    fn select_record(&self,
                     request: &str,
                     value: &str)
                     -> Receiver<Result<DomainRecord, DomainError>> {
        let (tx, rx) = channel();

        // Run the sql command on a pooled thread.
        let pool = self.pool.clone();
        let value = value.to_owned();
        let request = request.to_owned();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            let mut stmt = sqltry!(conn.prepare(&request), tx);
            let mut rows = sqltry!(stmt.query(&[&value]), tx);
            if let Some(result_row) = rows.next() {
                let row = sqltry!(result_row, tx);
                tx.send(Ok(DomainRecord::from_sql(row))).unwrap();
            } else {
                tx.send(Err(DomainError::NoRecord)).unwrap();
            }
        });

        rx
    }

    pub fn get_record_by_name(&self, name: &str) -> Receiver<Result<DomainRecord, DomainError>> {
        self.select_record("SELECT name, token, dns_challenge, local_ip, public_ip, timestamp \
                            FROM domains WHERE name=$1",
                           name)
    }

    pub fn get_record_by_token(&self, token: &str) -> Receiver<Result<DomainRecord, DomainError>> {
        self.select_record("SELECT name, token, dns_challenge, local_ip, public_ip, timestamp \
                            FROM domains WHERE token=$1",
                           token)
    }

    pub fn add_record(&self, record: DomainRecord) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let record = record.clone();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
            sqltry!(conn.execute("INSERT INTO domains VALUES ($1, $2, $3, $4, $5, $6)",
                                 &[&record.name,
                                   &record.token,
                                   &record.dns_challenge.unwrap_or("".to_owned()),
                                   &record.local_ip.unwrap_or("".to_owned()),
                                   &record.public_ip.unwrap_or("".to_owned()),
                                   &record.timestamp]),
                    tx);
            tx.send(Ok(())).unwrap();
        });

        rx
    }

    pub fn update_record(&self, record: DomainRecord) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let record = record.clone();
        thread::spawn(move || {
            let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);

            sqltry!(conn.execute("UPDATE domains SET dns_challenge=$1, local_ip=$2, public_ip=$3, timestamp=$4 \
                                  WHERE name=$5 AND token=$6",
                                 &[&record.dns_challenge.unwrap_or("".to_owned()),
                                   &record.local_ip.unwrap_or("".to_owned()),
                                   &record.public_ip.unwrap_or("".to_owned()),
                                   &record.timestamp,
                                   &record.name,
                                   &record.token]),
                    tx);
            tx.send(Ok(())).unwrap();
        });

        rx
    }

    // Returns the number of rows affected.
    pub fn execute_1param_sql(&self,
                              request: &str,
                              value: SqlParam)
                              -> Receiver<Result<i32, DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        let value = value.to_owned();
        let request = request.to_owned();
        thread::spawn(move || {
                          let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
                          let res = sqltry!(conn.execute(&request, &[&value]), tx);
                          tx.send(Ok(res)).unwrap();
                      });

        rx
    }

    pub fn delete_record_by_name(&self, name: &str) -> Receiver<Result<i32, DomainError>> {
        self.execute_1param_sql("DELETE FROM domains WHERE name=$1",
                                SqlParam::Text(name.to_owned()))
    }

    pub fn delete_record_by_token(&self, token: &str) -> Receiver<Result<i32, DomainError>> {
        self.execute_1param_sql("DELETE FROM domains WHERE token=$1",
                                SqlParam::Text(token.to_owned()))
    }

    pub fn flush(&self) -> Receiver<Result<(), DomainError>> {
        let (tx, rx) = channel();

        let pool = self.pool.clone();
        thread::spawn(move || {
                          let conn = sqltry!(pool.get(), tx, DomainError::DbUnavailable);
                          sqltry!(conn.execute("DELETE FROM domains", &[]), tx);
                          tx.send(Ok(())).unwrap();
                      });
        rx
    }
}

#[test]
fn test_domain_store() {
    let db = DomainDb::new("domain_db_test.sqlite");

    // Start with an empty db.
    db.flush().recv().unwrap().expect("Flushing the db");

    // Check that we don't find any record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        Err(err) => panic!("Should not find a record by name in an empty db. {:?}", err),
        _ => panic!("Should not find a record by name in an empty db."),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find a record by token in an empty db."),
    }

    // Add a record without a dns challenge.
    let no_challenge_record =
        DomainRecord::new("test.example.org", "test-token", None, None, None, 0);
    db.add_record(no_challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Adding the no_challenge record");

    // Check that we can find it and that it matches our record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Ok(record) => assert_eq!(record, no_challenge_record),
        Err(err) => panic!("Failed to find record by name: {:?}", err),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Ok(record) => assert_eq!(record, no_challenge_record),
        Err(err) => panic!("Failed to find record by token: {:?}", err),
    }

    // Update the record to have challenge.
    let challenge_record = DomainRecord::new("test.example.org",
                                             "test-token",
                                             Some("dns-challenge"),
                                             None,
                                             None,
                                             0);
    db.update_record(challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Updating the challenge record");

    // Check that we can find it and that it matches our record.
    match db.get_record_by_name("test.example.org")
              .recv()
              .unwrap() {
        Ok(record) => assert_eq!(record, challenge_record),
        Err(err) => panic!("Failed to find record by name: {:?}", err),
    }

    match db.get_record_by_token("test-token").recv().unwrap() {
        Ok(record) => assert_eq!(record, challenge_record),
        Err(err) => panic!("Failed to find record by token: {:?}", err),
    }

    // Remove by name.
    db.delete_record_by_name(&challenge_record.name)
        .recv()
        .unwrap()
        .expect("Should delete record by name");
    match db.get_record_by_name(&challenge_record.name)
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find this record anymore."),
    }

    // Add again, and remove by token.
    db.add_record(challenge_record.clone())
        .recv()
        .unwrap()
        .expect("Adding the challenge record");
    db.delete_record_by_token(&challenge_record.token)
        .recv()
        .unwrap()
        .expect("Should delete record by token");
    match db.get_record_by_name(&challenge_record.name)
              .recv()
              .unwrap() {
        Err(DomainError::NoRecord) => {}
        _ => panic!("Should not find this record anymore."),
    }
}
<|endoftext|>


<commit_before>use super::{lookups};

use cards::card::{Card};

use super::{HandRank};
use super::utils::{card_to_deck_number};
use super::types::{InternalCardRepresentation};

fn findit(key: usize) -> usize {
    let mut low = 0;
    let mut high = 4887;
    let mut mid;

    while low <= high {
        mid = (high+low) >> 1; // divide by two

        if key < lookups::PRODUCTS[mid] as usize {
            high = mid - 1;
        } else if key > lookups::PRODUCTS[mid] as usize {
            low = mid + 1;
        } else {
            return mid;
        }
    }

    panic!("No match found")
}

type ICR = InternalCardRepresentation;
pub fn eval_5cards_raw(cards: [&InternalCardRepresentation; 5]) -> HandRank {
    let c1 = cards[0];
    let c2 = cards[1];
    let c3 = cards[2];
    let c4 = cards[3];
    let c5 = cards[4];

    let q : usize = ((c1 | c2 | c3 | c4 | c5) as usize) >> 16;

    if (c1 & c2 & c3 & c4 & c5 & 0xf000) != 0 {
        return lookups::FLUSHES[q] as HandRank;
    }
    let s = lookups::UNIQUE_5[q] as HandRank;
    if s != 0 {
        return s;
    }

    let q =
        ((c1 & 0xff) * (c2 & 0xff) * (c3 & 0xff) * (c4 & 0xff) * (c5 & 0xff))
        as usize;
    let lookup = findit(q);
    lookups::VALUES[lookup] as HandRank
}

pub fn eval_5cards(cards: [&Card; 5]) -> HandRank {
    let c1 = card_to_deck_number(cards[0]);
    let c2 = card_to_deck_number(cards[1]);
    let c3 = card_to_deck_number(cards[2]);
    let c4 = card_to_deck_number(cards[3]);
    let c5 = card_to_deck_number(cards[4]);

    eval_5cards_raw([&c1, &c2, &c3, &c4, &c5])
}

pub fn eval_7cards(cards: [&Card; 7]) -> HandRank {
    let mut tmp;
    let mut best = 9999;
    for ids in lookups::PERM_7.iter() {
        let subhand : [&Card; 5] = [
                cards[ids[0] as usize],
                cards[ids[1] as usize],
                cards[ids[2] as usize],
                cards[ids[3] as usize],
                cards[ids[4] as usize]
         ];

        tmp = eval_5cards(subhand);
        if tmp < best {
            best = tmp;
        }
    }

    best
}

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

    use cards::deck::{Deck};
    use holdem::{HandRankClass};

    use super::super::{hand_rank}; //TODO this super::super business is kind of ugly
    use super::super::{HandRank};

    use super::{eval_5cards_raw};
    use super::super::utils::{card_to_deck_number};
    use super::super::types::{InternalCardRepresentation};

    // TODO: this is not really specific to this evaluation method. It could as well live in the library tests folder
    // the reason it is here, is due to the internal representation hackage which got ditched
    #[test]
    fn evaluate_all_possible_5_card_combinations() {
        let mut deck = Deck::new();
        let mut cards : [InternalCardRepresentation; 52] = [0; 52];

        // this could be made faster, by creating a function that works on raw-card-representations and translating
        // the deck cards into it
        for i in 0..52 {
            let card = deck.draw();
            cards[i] = card_to_deck_number(&card);
        }

        let mut rank_class_count : HashMap<HandRankClass, usize> = HashMap::new();
        let mut rank_count : HashMap<HandRank, bool> = HashMap::new();

        // 2,598,960 unique poker hands
        for i1 in 0..52 {
            for i2 in (i1+1)..52 {
                for i3 in (i2+1)..52 {
                    for i4 in (i3+1)..52 {
                        for i5 in (i4+1)..52 {
                            let c1 = &cards[i1];
                            let c2 = &cards[i2];
                            let c3 = &cards[i3];
                            let c4 = &cards[i4];
                            let c5 = &cards[i5];

                            let rank = eval_5cards_raw([c1, c2, c3, c4, c5]);
                            // mark the rank in the map
                            rank_count.entry(rank).or_insert(true);
                        }
                    }
                }
            }
        }

        // count distinct ranks for each rank class
        for key in rank_count.keys() {
            let rank_class = hand_rank(*key);

            let count = rank_class_count.entry(rank_class).or_insert(0);
            *count += 1;
        }

        assert_eq!(*rank_class_count.get(&HandRankClass::HighCard).unwrap(), 1277);
        assert_eq!(*rank_class_count.get(&HandRankClass::OnePair).unwrap(), 2860);
        assert_eq!(*rank_class_count.get(&HandRankClass::TwoPair).unwrap(), 858);
        assert_eq!(*rank_class_count.get(&HandRankClass::ThreeOfAKind).unwrap(), 858);
        assert_eq!(*rank_class_count.get(&HandRankClass::Straight).unwrap(), 10);
        assert_eq!(*rank_class_count.get(&HandRankClass::Flush).unwrap(), 1277);
        assert_eq!(*rank_class_count.get(&HandRankClass::FullHouse).unwrap(), 156);
        assert_eq!(*rank_class_count.get(&HandRankClass::FourOfAKind).unwrap(), 156);
        assert_eq!(*rank_class_count.get(&HandRankClass::StraightFlush).unwrap(), 10);

        // this is a bit redundant
        // there should be 7462 unique ranks, in accordance with the hand_rank function
        assert_eq!(rank_count.len(), 7462);
    }
}

<commit_msg>Test raw and original card evaluation.<commit_after>use super::{lookups};

use cards::card::{Card};

use super::{HandRank};
use super::utils::{card_to_deck_number};
use super::types::{InternalCardRepresentation};

fn findit(key: usize) -> usize {
    let mut low = 0;
    let mut high = 4887;
    let mut mid;

    while low <= high {
        mid = (high+low) >> 1; // divide by two

        if key < lookups::PRODUCTS[mid] as usize {
            high = mid - 1;
        } else if key > lookups::PRODUCTS[mid] as usize {
            low = mid + 1;
        } else {
            return mid;
        }
    }

    panic!("No match found")
}

type ICR = InternalCardRepresentation;
pub fn eval_5cards_raw(cards: [&InternalCardRepresentation; 5]) -> HandRank {
    let c1 = cards[0];
    let c2 = cards[1];
    let c3 = cards[2];
    let c4 = cards[3];
    let c5 = cards[4];

    let q : usize = ((c1 | c2 | c3 | c4 | c5) as usize) >> 16;

    if (c1 & c2 & c3 & c4 & c5 & 0xf000) != 0 {
        return lookups::FLUSHES[q] as HandRank;
    }
    let s = lookups::UNIQUE_5[q] as HandRank;
    if s != 0 {
        return s;
    }

    let q =
        ((c1 & 0xff) * (c2 & 0xff) * (c3 & 0xff) * (c4 & 0xff) * (c5 & 0xff))
        as usize;
    let lookup = findit(q);
    lookups::VALUES[lookup] as HandRank
}

pub fn eval_5cards(cards: [&Card; 5]) -> HandRank {
    let c1 = card_to_deck_number(cards[0]);
    let c2 = card_to_deck_number(cards[1]);
    let c3 = card_to_deck_number(cards[2]);
    let c4 = card_to_deck_number(cards[3]);
    let c5 = card_to_deck_number(cards[4]);

    eval_5cards_raw([&c1, &c2, &c3, &c4, &c5])
}

pub fn eval_7cards(cards: [&Card; 7]) -> HandRank {
    let mut tmp;
    let mut best = 9999;
    for ids in lookups::PERM_7.iter() {
        let subhand : [&Card; 5] = [
                cards[ids[0] as usize],
                cards[ids[1] as usize],
                cards[ids[2] as usize],
                cards[ids[3] as usize],
                cards[ids[4] as usize]
         ];

        tmp = eval_5cards(subhand);
        if tmp < best {
            best = tmp;
        }
    }

    best
}

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

    use cards::deck::{Deck};
    use cards::card::{Card};
    use holdem::{HandRankClass};

    use super::super::{hand_rank}; //TODO this super::super business is kind of ugly
    use super::super::{HandRank};

    use super::{eval_5cards_raw,eval_5cards};
    use super::super::utils::{card_to_deck_number};
    use super::super::types::{InternalCardRepresentation};

    // TODO: this is not really specific to this evaluation method. It could as well live in the library tests folder
    // the reason it is here, is due to the internal representation hackage which got ditched
    #[test]
    fn evaluate_all_possible_5_card_combinations_quick() {
        let mut deck = Deck::new();
        let mut cards : [InternalCardRepresentation; 52] = [0; 52];

        // this could be made faster, by creating a function that works on raw-card-representations and translating
        // the deck cards into it
        for i in 0..52 {
            let card = deck.draw();
            cards[i] = card_to_deck_number(&card);
        }

        let mut rank_class_count : HashMap<HandRankClass, usize> = HashMap::new();
        let mut rank_count : HashMap<HandRank, bool> = HashMap::new();

        // 2,598,960 unique poker hands
        for i1 in 0..52 {
            for i2 in (i1+1)..52 {
                for i3 in (i2+1)..52 {
                    for i4 in (i3+1)..52 {
                        for i5 in (i4+1)..52 {
                            let c1 = &cards[i1];
                            let c2 = &cards[i2];
                            let c3 = &cards[i3];
                            let c4 = &cards[i4];
                            let c5 = &cards[i5];

                            let rank = eval_5cards_raw([c1, c2, c3, c4, c5]);
                            // mark the rank in the map
                            rank_count.entry(rank).or_insert(true);
                        }
                    }
                }
            }
        }

        // count distinct ranks for each rank class
        for key in rank_count.keys() {
            let rank_class = hand_rank(*key);

            let count = rank_class_count.entry(rank_class).or_insert(0);
            *count += 1;
        }

        assert_eq!(*rank_class_count.get(&HandRankClass::HighCard).unwrap(), 1277);
        assert_eq!(*rank_class_count.get(&HandRankClass::OnePair).unwrap(), 2860);
        assert_eq!(*rank_class_count.get(&HandRankClass::TwoPair).unwrap(), 858);
        assert_eq!(*rank_class_count.get(&HandRankClass::ThreeOfAKind).unwrap(), 858);
        assert_eq!(*rank_class_count.get(&HandRankClass::Straight).unwrap(), 10);
        assert_eq!(*rank_class_count.get(&HandRankClass::Flush).unwrap(), 1277);
        assert_eq!(*rank_class_count.get(&HandRankClass::FullHouse).unwrap(), 156);
        assert_eq!(*rank_class_count.get(&HandRankClass::FourOfAKind).unwrap(), 156);
        assert_eq!(*rank_class_count.get(&HandRankClass::StraightFlush).unwrap(), 10);

        // this is a bit redundant
        // there should be 7462 unique ranks, in accordance with the hand_rank function
        assert_eq!(rank_count.len(), 7462);
    }

    #[test]
    fn evaluate_all_possible_5_card_combinations() {
        let mut deck = Deck::new();
        let mut cards : Vec<Card> = Vec::new();

        for _ in 0..52 {
            let card = deck.draw();
            cards.push(card);
        }

        let mut rank_class_count : HashMap<HandRankClass, usize> = HashMap::new();
        let mut rank_count : HashMap<HandRank, bool> = HashMap::new();

        // 2,598,960 unique poker hands
        for i1 in 0..52 {
            for i2 in (i1+1)..52 {
                for i3 in (i2+1)..52 {
                    for i4 in (i3+1)..52 {
                        for i5 in (i4+1)..52 {
                            let c1 = &cards[i1];
                            let c2 = &cards[i2];
                            let c3 = &cards[i3];
                            let c4 = &cards[i4];
                            let c5 = &cards[i5];

                            let rank = eval_5cards([c1, c2, c3, c4, c5]);
                            // mark the rank in the map
                            rank_count.entry(rank).or_insert(true);
                        }
                    }
                }
            }
        }

        // count distinct ranks for each rank class
        for key in rank_count.keys() {
            let rank_class = hand_rank(*key);

            let count = rank_class_count.entry(rank_class).or_insert(0);
            *count += 1;
        }

        assert_eq!(*rank_class_count.get(&HandRankClass::HighCard).unwrap(), 1277);
        assert_eq!(*rank_class_count.get(&HandRankClass::OnePair).unwrap(), 2860);
        assert_eq!(*rank_class_count.get(&HandRankClass::TwoPair).unwrap(), 858);
        assert_eq!(*rank_class_count.get(&HandRankClass::ThreeOfAKind).unwrap(), 858);
        assert_eq!(*rank_class_count.get(&HandRankClass::Straight).unwrap(), 10);
        assert_eq!(*rank_class_count.get(&HandRankClass::Flush).unwrap(), 1277);
        assert_eq!(*rank_class_count.get(&HandRankClass::FullHouse).unwrap(), 156);
        assert_eq!(*rank_class_count.get(&HandRankClass::FourOfAKind).unwrap(), 156);
        assert_eq!(*rank_class_count.get(&HandRankClass::StraightFlush).unwrap(), 10);

        // this is a bit redundant
        // there should be 7462 unique ranks, in accordance with the hand_rank function
        assert_eq!(rank_count.len(), 7462);
    }
}

<|endoftext|>


<commit_before>// Copyright (C) 2019 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

use actix::Addr;
use slog_scope::{debug, error, info};
use std::sync::mpsc;

/// [Controller] is used to [start](Controller::start),
/// [stop](Controller::stop) and [restart](Controller::restart) the stepper
/// thread (see [start](Controller::start)).
#[derive(Debug, Default)]
pub(crate) struct Controller {
    terminate: Option<mpsc::Sender<TerminateThread>>,
}

/// Used on channel communication between the stepper thread and the
/// [Controller](Controller) to request that the stepper to stop.
struct TerminateThread;

impl Controller {
    /// Stops the current running stepper (if any) and starts a new one for the
    /// supplied Actor.
    pub(super) fn restart<A>(&mut self, addr: Addr<A>)
    where
        A: actix::Handler<super::Step>,
        A::Context: actix::dev::ToEnvelope<A, super::Step>,
    {
        info!("Restarting stepper");
        self.stop();
        self.start(addr);
    }

    /// Stops the stepper if it's currently running.
    pub(super) fn stop(&mut self) {
        info!("Stopping stepper");
        if let Some(sndr) = self.terminate.take() {
            // send mpsc::Sender::send Err means the channel is closed and thus the
            // thread has stopped already
            let _ = sndr.send(TerminateThread);
        }
    }

    /// Starts the stepper for the supplied Actor's address.
    ///
    /// The stepper is a thread that sends [super::Step] messages
    /// to the  actor until the step is replayed with a
    /// [super::StepTransition::Never], or a [TerminateThread] message
    /// is received.
    pub(super) fn start<A>(&mut self, addr: Addr<A>)
    where
        A: actix::Handler<super::Step>,
        A::Context: actix::dev::ToEnvelope<A, super::Step>,
    {
        info!("Starting stepper");
        let (sndr, recv) = mpsc::channel();
        self.terminate = Some(sndr);

        // We ignore errors raised by the stepper
        let arbiter = actix::Arbiter::new();

        arbiter.exec_fn(|| {
            actix::Arbiter::spawn(async move {
                while recv.try_recv().is_err() {
                    match addr.send(super::Step).await {
                        Err(e) => {
                            error!("Communication to actor failed: {:?}", e);
                        }
                        Ok(super::StepTransition::Immediate) => {}
                        Ok(super::StepTransition::Delayed(t)) => {
                            debug!("Sleeping stepper thread for: {} seconds", t.as_secs());
                            actix::clock::delay_for(t).await;
                        }
                        Ok(super::StepTransition::Never) => {
                            info!("Stopping step messages");
                            break;
                        }
                    }
                }
            })
        });
    }
}

impl Drop for Controller {
    fn drop(&mut self) {
        self.stop();
    }
}
<commit_msg>updatehub: Never restart stepper on tests<commit_after>// Copyright (C) 2019 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

use actix::Addr;
use slog_scope::{debug, error, info};
use std::sync::mpsc;

/// [Controller] is used to [start](Controller::start),
/// [stop](Controller::stop) and [restart](Controller::restart) the stepper
/// thread (see [start](Controller::start)).
#[derive(Debug, Default)]
pub(crate) struct Controller {
    terminate: Option<mpsc::Sender<TerminateThread>>,
}

/// Used on channel communication between the stepper thread and the
/// [Controller](Controller) to request that the stepper to stop.
struct TerminateThread;

impl Controller {
    /// Stops the current running stepper (if any) and starts a new one for the
    /// supplied Actor.
    #[cfg(not(test))]
    pub(super) fn restart<A>(&mut self, addr: Addr<A>)
    where
        A: actix::Handler<super::Step>,
        A::Context: actix::dev::ToEnvelope<A, super::Step>,
    {
        info!("Restarting stepper");
        self.stop();
        self.start(addr);
    }

    /// Most actor's tests relay on steppter not running, this should
    /// ensure the stepper doesn't get restarted automatically.
    #[cfg(test)]
    pub(super) fn restart<A>(&mut self, _: Addr<A>)
    where
        A: actix::Handler<super::Step>,
        A::Context: actix::dev::ToEnvelope<A, super::Step>,
    {
    }

    /// Stops the stepper if it's currently running.
    pub(super) fn stop(&mut self) {
        info!("Stopping stepper");
        if let Some(sndr) = self.terminate.take() {
            // send mpsc::Sender::send Err means the channel is closed and thus the
            // thread has stopped already
            let _ = sndr.send(TerminateThread);
        }
    }

    /// Starts the stepper for the supplied Actor's address.
    ///
    /// The stepper is a thread that sends [super::Step] messages
    /// to the  actor until the step is replayed with a
    /// [super::StepTransition::Never], or a [TerminateThread] message
    /// is received.
    pub(super) fn start<A>(&mut self, addr: Addr<A>)
    where
        A: actix::Handler<super::Step>,
        A::Context: actix::dev::ToEnvelope<A, super::Step>,
    {
        info!("Starting stepper");
        let (sndr, recv) = mpsc::channel();
        self.terminate = Some(sndr);

        // We ignore errors raised by the stepper
        let arbiter = actix::Arbiter::new();

        arbiter.exec_fn(|| {
            actix::Arbiter::spawn(async move {
                while recv.try_recv().is_err() {
                    match addr.send(super::Step).await {
                        Err(e) => {
                            error!("Communication to actor failed: {:?}", e);
                        }
                        Ok(super::StepTransition::Immediate) => {}
                        Ok(super::StepTransition::Delayed(t)) => {
                            debug!("Sleeping stepper thread for: {} seconds", t.as_secs());
                            actix::clock::delay_for(t).await;
                        }
                        Ok(super::StepTransition::Never) => {
                            info!("Stopping step messages");
                            break;
                        }
                    }
                }
            })
        });
    }
}

impl Drop for Controller {
    fn drop(&mut self) {
        self.stop();
    }
}
<|endoftext|>


<commit_before>use helpers::*;
use basename;
use chop_basename;
use dirname;
use extname;
use plus;

use ruru;
use ruru::{RString, Boolean, Array};
use std::path::{MAIN_SEPARATOR,Path};
use std::fs;

type MaybeString = Result<ruru::RString, ruru::result::Error>;
type MaybeBoolean = Result<ruru::Boolean, ruru::result::Error>;

pub fn pn_add_trailing_separator(pth: MaybeString) -> RString {
  let p = pth.ok().unwrap();
  let x = format!("{}{}", p.to_str(), "a");
  match x.rsplit_terminator(MAIN_SEPARATOR).next() {
    Some("a") => p,
    _ => RString::new(format!("{}{}", p.to_str(), MAIN_SEPARATOR).as_str())
  }
}

pub fn pn_is_absolute(pth: MaybeString) -> Boolean {
  Boolean::new(match pth.ok().unwrap_or(RString::new("")).to_str().chars().next() {
    Some(c) => c == MAIN_SEPARATOR,
    None => false
  })
}

// pub fn pn_ascend(){}

pub fn pn_basename(pth: MaybeString, ext: MaybeString) -> RString {
  RString::new(
    basename::basename(
      pth.ok().unwrap_or(RString::new("")).to_str(),
      ext.ok().unwrap_or(RString::new("")).to_str()
    )
  )
}

pub fn pn_children(pth: MaybeString, with_dir: MaybeBoolean) -> Array {
  let rstring = pth.ok().unwrap_or(RString::new("."));
  let val = rstring.to_str();
  let mut with_directory = with_dir.ok().unwrap_or(Boolean::new(true)).to_bool();
  if val == "." {
    with_directory = false;
  }
  let mut arr = Array::new();

  if let Ok(entries) = fs::read_dir(val) {
    for entry in entries {
      if with_directory {
        match entry {
          Ok(v) => { arr.push(RString::new(v.path().to_str().unwrap())); },
          _ => {}
        };
      } else {
        match entry {
          Ok(v) => { arr.push(RString::new(v.file_name().to_str().unwrap())); },
          _ => {}
        };
      }
    }
  }
  arr
}

pub fn pn_children_compat(pth: MaybeString, with_dir: MaybeBoolean) -> Array {
  let rstring = pth.ok().unwrap_or(RString::new("."));
  let val = rstring.to_str();
  let mut with_directory = with_dir.ok().unwrap_or(Boolean::new(true)).to_bool();
  if val == "." {
    with_directory = false;
  }
  let mut arr = Array::new();

  if let Ok(entries) = fs::read_dir(val) {
    for entry in entries {
      if with_directory {
        match entry {
          Ok(v) => { arr.push(
              class_new("Pathname", vec![str_to_any_obj(v.path().to_str().unwrap())])
            );
          },
          _ => {}
        };
      } else {
        match entry {
          Ok(v) => { arr.push(
              class_new("Pathname", vec![str_to_any_obj(v.file_name().to_str().unwrap())])
            );
          },
          _ => {}
        };
      }
    }
  }
  arr
}

pub fn pn_chop_basename(pth: MaybeString) -> Array {
  let mut arr = Array::with_capacity(2);
  let results = chop_basename::chop_basename(pth.ok().unwrap_or(RString::new("")).to_str());
  match results {
    Some((dirname, basename)) => {
      arr.push(RString::new(&dirname[..]));
      arr.push(RString::new(&basename[..]));
      arr
    },
    None => arr
  }
}

// pub fn pn_cleanpath(pth: MaybeString){}

// pub fn pn_cleanpath_aggressive(pth: MaybeString){}

// pub fn pn_cleanpath_conservative(pth: MaybeString){}

// pub fn pn_del_trailing_separator(pth: MaybeString){}

// pub fn pn_descend(){}

pub fn pn_is_directory(pth: MaybeString) -> Boolean {
  Boolean::new(
    Path::new(
      pth.ok().unwrap_or(RString::new("")).to_str()
    ).is_dir()
  )
}

pub fn pn_dirname(pth: MaybeString) -> RString {
  RString::new(
    &dirname::dirname(
      pth.ok().unwrap_or(RString::new("")).to_str()
    )[..]
  )
}

// pub fn pn_each_child(){}

// pub fn pn_each_filename(pth: MaybeString) -> NilClass {
//   NilClass::new()
// }

pub fn pn_entries(pth: MaybeString) -> Array {
  let files = fs::read_dir(pth.ok().unwrap_or(RString::new("")).to_str()).unwrap();
  let mut arr = Array::new();

  arr.push(RString::new("."));
  arr.push(RString::new(".."));

  for file in files {
    let file_name_str = file.unwrap().file_name().into_string().unwrap();
    arr.push(RString::new(&file_name_str[..]));
  }

  arr
}

pub fn pn_entries_compat(pth: MaybeString) -> Array {
  let files = fs::read_dir(pth.ok().unwrap_or(RString::new("")).to_str()).unwrap();
  let mut arr = Array::new();

  arr.push(class_new("Pathname", vec![str_to_any_obj(&"."[..])]));
  arr.push(class_new("Pathname", vec![str_to_any_obj(&".."[..])]));

  for file in files {
    let file_name_str = file.unwrap().file_name().into_string().unwrap();
    arr.push(class_new("Pathname", vec![str_to_any_obj(&file_name_str[..])]));
  }

  arr
}

pub fn pn_extname(pth: MaybeString) -> RString {
  RString::new(
    extname::extname(pth.ok().unwrap_or(RString::new("")).to_str())
  )
}

// pub fn pn_find(pth: MaybeString ,ignore_error: Boolean){}

pub fn pn_has_trailing_separator(pth: MaybeString) -> Boolean {
  let v = pth.ok().unwrap_or(RString::new(""));
  match chop_basename::chop_basename(v.to_str()) {
    Some((a,b)) => {
      Boolean::new(a.len() + b.len() < v.to_str().len())
    },
    _ => Boolean::new(false)
  }
}

// pub fn pn_join(args: Array){}

// pub fn pn_mkpath(pth: MaybeString) -> NilClass {
//   NilClass::new()
// }

// pub fn pn_is_mountpoint(pth: MaybeString){}

// pub fn pn_parent(pth: MaybeString){}

// also need impl +
pub fn pn_plus(pth1: MaybeString, pth2: MaybeString) -> RString {
  RString::new(&plus::plus_paths(pth1.ok().unwrap().to_str(), pth2.ok().unwrap().to_str())[..])
}

// pub fn pn_prepend_prefix(prefix: MaybeString, relpath: MaybeString){}

pub fn pn_is_relative(pth: MaybeString) -> Boolean {
  Boolean::new(
    match pth.ok().unwrap_or(RString::new(&MAIN_SEPARATOR.to_string()[..])).to_str().chars().next() {
      Some(c) => c != MAIN_SEPARATOR,
      None => true
    }
  )
}

// pub fn pn_root(pth: MaybeString){}

// pub fn pn_split_names(pth: MaybeString){}

// pub fn pn_relative_path_from(){}

// pub fn pn_rmtree(pth: MaybeString) -> NilClass {
//   NilClass::new()
// }
<commit_msg>remove unnecessary string conversion from pn_basename<commit_after>use helpers::*;
use basename;
use chop_basename;
use dirname;
use extname;
use plus;

use ruru;
use ruru::{RString, Boolean, Array};
use std::path::{MAIN_SEPARATOR,Path};
use std::fs;

type MaybeString = Result<ruru::RString, ruru::result::Error>;
type MaybeBoolean = Result<ruru::Boolean, ruru::result::Error>;

pub fn pn_add_trailing_separator(pth: MaybeString) -> RString {
  let p = pth.ok().unwrap();
  let x = format!("{}{}", p.to_str(), "a");
  match x.rsplit_terminator(MAIN_SEPARATOR).next() {
    Some("a") => p,
    _ => RString::new(format!("{}{}", p.to_str(), MAIN_SEPARATOR).as_str())
  }
}

pub fn pn_is_absolute(pth: MaybeString) -> Boolean {
  Boolean::new(match pth.ok().unwrap_or(RString::new("")).to_str().chars().next() {
    Some(c) => c == MAIN_SEPARATOR,
    None => false
  })
}

// pub fn pn_ascend(){}

pub fn pn_basename(pth: MaybeString, ext: MaybeString) -> RString {
  RString::new(
    basename::basename(
      pth.ok().unwrap_or(RString::new("")).to_str(),
      ext.ok().unwrap_or(RString::new("")).to_str()
    )
  )
}

pub fn pn_children(pth: MaybeString, with_dir: MaybeBoolean) -> Array {
  let rstring = pth.ok().unwrap_or(RString::new("."));
  let val = rstring.to_str();
  let mut with_directory = with_dir.ok().unwrap_or(Boolean::new(true)).to_bool();
  if val == "." {
    with_directory = false;
  }
  let mut arr = Array::new();

  if let Ok(entries) = fs::read_dir(val) {
    for entry in entries {
      if with_directory {
        match entry {
          Ok(v) => { arr.push(RString::new(v.path().to_str().unwrap())); },
          _ => {}
        };
      } else {
        match entry {
          Ok(v) => { arr.push(RString::new(v.file_name().to_str().unwrap())); },
          _ => {}
        };
      }
    }
  }
  arr
}

pub fn pn_children_compat(pth: MaybeString, with_dir: MaybeBoolean) -> Array {
  let rstring = pth.ok().unwrap_or(RString::new("."));
  let val = rstring.to_str();
  let mut with_directory = with_dir.ok().unwrap_or(Boolean::new(true)).to_bool();
  if val == "." {
    with_directory = false;
  }
  let mut arr = Array::new();

  if let Ok(entries) = fs::read_dir(val) {
    for entry in entries {
      if with_directory {
        match entry {
          Ok(v) => { arr.push(
              class_new("Pathname", vec![str_to_any_obj(v.path().to_str().unwrap())])
            );
          },
          _ => {}
        };
      } else {
        match entry {
          Ok(v) => { arr.push(
              class_new("Pathname", vec![str_to_any_obj(v.file_name().to_str().unwrap())])
            );
          },
          _ => {}
        };
      }
    }
  }
  arr
}

pub fn pn_chop_basename(pth: MaybeString) -> Array {
  let mut arr = Array::with_capacity(2);
  let results = chop_basename::chop_basename(pth.ok().unwrap_or(RString::new("")).to_str());
  match results {
    Some((dirname, basename)) => {
      arr.push(RString::new(&dirname[..]));
      arr.push(RString::new(&basename[..]));
      arr
    },
    None => arr
  }
}

// pub fn pn_cleanpath(pth: MaybeString){}

// pub fn pn_cleanpath_aggressive(pth: MaybeString){}

// pub fn pn_cleanpath_conservative(pth: MaybeString){}

// pub fn pn_del_trailing_separator(pth: MaybeString){}

// pub fn pn_descend(){}

pub fn pn_is_directory(pth: MaybeString) -> Boolean {
  Boolean::new(
    Path::new(
      pth.ok().unwrap_or(RString::new("")).to_str()
    ).is_dir()
  )
}

pub fn pn_dirname(pth: MaybeString) -> RString {
  RString::new(
    dirname::dirname(
      pth.ok().unwrap_or(RString::new("")).to_str()
    )
  )
}

// pub fn pn_each_child(){}

// pub fn pn_each_filename(pth: MaybeString) -> NilClass {
//   NilClass::new()
// }

pub fn pn_entries(pth: MaybeString) -> Array {
  let files = fs::read_dir(pth.ok().unwrap_or(RString::new("")).to_str()).unwrap();
  let mut arr = Array::new();

  arr.push(RString::new("."));
  arr.push(RString::new(".."));

  for file in files {
    let file_name_str = file.unwrap().file_name().into_string().unwrap();
    arr.push(RString::new(&file_name_str[..]));
  }

  arr
}

pub fn pn_entries_compat(pth: MaybeString) -> Array {
  let files = fs::read_dir(pth.ok().unwrap_or(RString::new("")).to_str()).unwrap();
  let mut arr = Array::new();

  arr.push(class_new("Pathname", vec![str_to_any_obj(&"."[..])]));
  arr.push(class_new("Pathname", vec![str_to_any_obj(&".."[..])]));

  for file in files {
    let file_name_str = file.unwrap().file_name().into_string().unwrap();
    arr.push(class_new("Pathname", vec![str_to_any_obj(&file_name_str[..])]));
  }

  arr
}

pub fn pn_extname(pth: MaybeString) -> RString {
  RString::new(
    extname::extname(pth.ok().unwrap_or(RString::new("")).to_str())
  )
}

// pub fn pn_find(pth: MaybeString ,ignore_error: Boolean){}

pub fn pn_has_trailing_separator(pth: MaybeString) -> Boolean {
  let v = pth.ok().unwrap_or(RString::new(""));
  match chop_basename::chop_basename(v.to_str()) {
    Some((a,b)) => {
      Boolean::new(a.len() + b.len() < v.to_str().len())
    },
    _ => Boolean::new(false)
  }
}

// pub fn pn_join(args: Array){}

// pub fn pn_mkpath(pth: MaybeString) -> NilClass {
//   NilClass::new()
// }

// pub fn pn_is_mountpoint(pth: MaybeString){}

// pub fn pn_parent(pth: MaybeString){}

// also need impl +
pub fn pn_plus(pth1: MaybeString, pth2: MaybeString) -> RString {
  RString::new(&plus::plus_paths(pth1.ok().unwrap().to_str(), pth2.ok().unwrap().to_str())[..])
}

// pub fn pn_prepend_prefix(prefix: MaybeString, relpath: MaybeString){}

pub fn pn_is_relative(pth: MaybeString) -> Boolean {
  Boolean::new(
    match pth.ok().unwrap_or(RString::new(&MAIN_SEPARATOR.to_string()[..])).to_str().chars().next() {
      Some(c) => c != MAIN_SEPARATOR,
      None => true
    }
  )
}

// pub fn pn_root(pth: MaybeString){}

// pub fn pn_split_names(pth: MaybeString){}

// pub fn pn_relative_path_from(){}

// pub fn pn_rmtree(pth: MaybeString) -> NilClass {
//   NilClass::new()
// }
<|endoftext|>


<commit_before>use std::io;
use std::iter::repeat;
use std::borrow::Cow;
use std::cell::RefCell;
use std::time::{Duration, Instant};
use std::sync::mpsc::{channel, Sender, Receiver};

use parking_lot::RwLock;

use term::{Term, terminal_size};
use utils::expand_template;

/// Controls the rendering style of progress bars.
pub struct Style {
    pub tick_chars: Vec<char>,
    pub progress_chars: Vec<char>,
    pub bar_template: Cow<'static, str>,
    pub spinner_template: Cow<'static, str>,
}

/// The drawn state of an element.
#[derive(Clone)]
pub struct DrawState {
    /// The lines to print (can contain ANSI codes)
    pub lines: Vec<String>,
    /// True if the bar no longer needs drawing.
    pub finished: bool,
    /// Time when the draw state was created.
    pub ts: Instant,
}

enum Status {
    InProgress,
    DoneVisible,
    DoneHidden,
}

/// Target for draw operations
pub enum DrawTarget {
    /// Draws into a terminal
    Term(Term, Option<DrawState>, Option<Duration>),
    /// Draws to a remote receiver
    Remote(usize, Sender<(usize, DrawState)>),
    /// Do not draw at all
    Hidden,
}

impl DrawTarget {
    /// Draw to a buffered stdout terminal
    pub fn stdout() -> DrawTarget {
        let rate = Some(Duration::from_millis(1000 / 30));
        DrawTarget::Term(Term::buffered_stdout(), None, rate)
    }

    /// Draw to a buffered stderr terminal
    pub fn stderr() -> DrawTarget {
        let rate = Some(Duration::from_millis(1000 / 30));
        DrawTarget::Term(Term::buffered_stderr(), None, rate)
    }

    /// Apply the given draw state (draws it).
    pub fn apply_draw_state(&mut self, draw_state: DrawState) -> io::Result<()> {
        match *self {
            DrawTarget::Term(ref term, ref mut last_state, rate) => {
                let last_draw = last_state.as_ref().map(|x| x.ts);
                if draw_state.finished ||
                   rate.is_none() ||
                   last_draw.is_none() ||
                   last_draw.unwrap().elapsed() > rate.unwrap() {
                    if let Some(ref last_state) = *last_state {
                        last_state.clear_term(term)?;
                    }
                    draw_state.draw_to_term(term)?;
                    term.flush()?;
                    *last_state = Some(draw_state);
                }
            }
            DrawTarget::Remote(idx, ref chan) => {
                chan.send((idx, draw_state)).unwrap();
            }
            DrawTarget::Hidden => {}
        }
        Ok(())
    }
}

impl DrawState {
    pub fn clear_term(&self, term: &Term) -> io::Result<()> {
        term.clear_last_lines(self.lines.len())
    }

    pub fn draw_to_term(&self, term: &Term) -> io::Result<()> {
        for line in &self.lines {
            term.write_line(line)?;
        }
        Ok(())
    }
}

impl Default for Style {
    fn default() -> Style {
        Style {
            tick_chars: "⠁⠁⠉⠙⠚⠒⠂⠂⠒⠲⠴⠤⠄⠄⠤⠠⠠⠤⠦⠖⠒⠐⠐⠒⠓⠋⠉⠈⠈ ".chars().collect(),
            progress_chars: "██░".chars().collect(),
            bar_template: Cow::Borrowed("{msg}\n{wide_bar} {pos}/{len}"),
            spinner_template: Cow::Borrowed("{spinner} {msg}"),
        }
    }
}

impl Style {

    /// Returns the tick char for a given number.
    pub fn get_tick_char(&self, idx: u64) -> char {
        self.tick_chars[(idx as usize) % (self.tick_chars.len() - 1)]
    }

    /// Returns the tick char for the finished state.
    pub fn get_final_tick_char(&self) -> char {
        self.tick_chars[self.tick_chars.len() - 1]
    }

    pub fn format_bar(&self, state: &ProgressState, width: usize) -> String {
        let pct = state.percent();
        let fill = (pct * width as f32) as usize;
        let bar = repeat(state.style.progress_chars[0]).take(fill).collect::<String>();
        let rest = repeat(state.style.progress_chars[2]).take(width - fill).collect::<String>();
        format!("{}{}", bar, rest)
    }

    pub fn format_state(&self, state: &ProgressState) -> Vec<String> {
        let (pos, len) = state.position();
        let mut rv = vec![];

        for line in self.bar_template.lines() {
            let need_wide_bar = RefCell::new(false);

            let s = expand_template(line, |key| {
                if key == "wide_bar" {
                    *need_wide_bar.borrow_mut() = true;
                    "\x00".into()
                } else if key == "bar" {
                    // XXX: width?
                    self.format_bar(state, 20)
                } else if key == "msg" {
                    state.message().to_string()
                } else if key == "pos" {
                    pos.to_string()
                } else if key == "len" {
                    len.to_string()
                } else {
                    "".into()
                }
            });

            rv.push(if *need_wide_bar.borrow() {
                let total_width = state.width();
                let bar_width = total_width - s.len() - 1;
                s.replace("\x00", &self.format_bar(state, bar_width))
            } else {
                s.to_string()
            });
        }

        rv
    }
}

/// The state of a progress bar at a moment in time.
pub struct ProgressState {
    style: Style,
    draw_target: DrawTarget,
    width: Option<u16>,
    message: String,
    pos: u64,
    len: u64,
    tick: u64,
    status: Status,
}

impl ProgressState {
    /// Returns the character that should be drawn for the
    /// current spinner character.
    pub fn current_tick_char(&self) -> char {
        if self.is_finished() {
            self.style.get_final_tick_char()
        } else {
            self.style.get_tick_char(self.tick)
        }
    }

    /// Indicates that a spinner should be drawn.
    pub fn has_spinner(&self) -> bool {
        self.tick != !0
    }

    /// Indicates that a progress bar should be drawn.
    pub fn has_progress(&self) -> bool {
        self.len != !0
    }

    /// Indicates that the progress bar finished.
    pub fn is_finished(&self) -> bool {
        match self.status {
            Status::InProgress => false,
            Status::DoneVisible => true,
            Status::DoneHidden => true,
        }
    }

    /// Returns `false` if the progress bar should no longer be
    /// drawn.
    pub fn should_render(&self) -> bool {
        match self.status {
            Status::DoneHidden => false,
            _ => true,
        }
    }

    /// Returns the completion in percent
    pub fn percent(&self) -> f32 {
        if self.len == !0 {
            0.0
        } else {
            self.pos as f32 / self.len as f32
        }
    }

    /// Returns the position of the status bar as `(pos, len)` tuple.
    pub fn position(&self) -> (u64, u64) {
        (self.pos, self.len)
    }

    /// Returns the current message of the progress bar.
    pub fn message(&self) -> &str {
        &self.message
    }

    /// The entire draw width
    pub fn width(&self) -> usize {
        if let Some(width) = self.width {
            width as usize
        } else {
            terminal_size().map(|x| x.1).unwrap_or(74) as usize
        }
    }
}

/// A progress bar or spinner.
pub struct ProgressBar {
    state: RwLock<ProgressState>,
}

impl ProgressBar {
    /// Creates a new progress bar with a given length.
    ///
    /// This progress bar by default draws directly to stdout.
    pub fn new(len: u64) -> ProgressBar {
        ProgressBar {
            state: RwLock::new(ProgressState {
                style: Default::default(),
                draw_target: DrawTarget::stdout(),
                width: None,
                message: "".into(),
                pos: 0,
                len: len,
                tick: !0,
                status: Status::InProgress,
            }),
        }
    }

    /// Creates a new spinner.
    ///
    /// This spinner by default draws directly to stdout.
    pub fn new_spinner() -> ProgressBar {
        let rv = ProgressBar::new(!0);
        rv.enable_spinner();
        rv
    }

    /// Enables the spinner.
    ///
    /// This is obviously enabled by default if you create a progress
    /// bar with `new_spinner` but optionally a spinner can be added
    /// to a progress bar itself.
    pub fn enable_spinner(&self) {
        self.update_state(|mut state| {
            if state.tick == !0 {
                state.tick = 0;
            }
        });
    }

    /// Disables a spinner.
    ///
    /// This should not be called if the progress bar is a spinner itself.
    pub fn disable_spinner(&self) {
        self.update_state(|mut state| {
            if state.tick != !0 {
                state.tick = !0;
            }
        });
    }

    /// Manually ticks the spinner or progress bar.
    ///
    /// This automatically happens on any other change to a progress bar.
    pub fn tick(&self) {
        self.update_state(|mut state| {
            if state.tick == !0 {
                state.tick = 0;
            } else {
                state.tick += 1;
            }
        });
    }

    /// Advances the position of a progress bar by delta.
    pub fn inc(&self, delta: u64) {
        self.update_state(|mut state| {
            state.pos += delta;
            if state.tick != !0 {
                state.tick += 1;
            }
        })
    }

    /// Sets the length of the progress bar.
    pub fn set_length(&self, len: u64) {
        self.update_state(|mut state| {
            state.len = len;
        })
    }

    /// Sets the current message of the progress bar.
    pub fn set_message(&self, msg: &str) {
        let msg = msg.to_string();
        self.update_state(|mut state| {
            state.message = msg;
        })
    }

    /// Finishes the progress bar and sets a message.
    pub fn finish_with_message(&self, msg: &str) {
        let msg = msg.to_string();
        self.update_state(|mut state| {
            state.message = msg;
            state.pos = state.len;
            state.status = Status::DoneVisible;
        });
    }

    /// Finishes the progress bar and completely clears it.
    pub fn finish_and_clear(&self) {
        self.update_state(|mut state| {
            state.pos = state.len;
            state.status = Status::DoneHidden;
        });
    }

    /// Sets a different draw target for the progress bar.
    ///
    /// This can be used to draw the progress bar to stderr
    /// for instance:
    ///
    /// ```rust,no_run
    /// # use indicatif::{ProgressBar, DrawTarget};
    /// let pb = ProgressBar::new(100);
    /// pb.set_draw_target(DrawTarget::stderr());
    /// ```
    pub fn set_draw_target(&self, target: DrawTarget) {
        self.state.write().draw_target = target;
    }

    fn update_state<F: FnOnce(&mut ProgressState)>(&self, f: F) {
        {
            let mut state = self.state.write();
            f(&mut state);
        }
        self.draw().ok();
    }

    fn draw(&self) -> io::Result<()> {
        let mut state = self.state.write();
        let draw_state = DrawState {
            lines: if state.should_render() {
                state.style.format_state(&*state)
            } else {
                vec![]
            },
            finished: state.is_finished(),
            ts: Instant::now(),
        };
        state.draw_target.apply_draw_state(draw_state)
    }
}

/// Manages multiple progress bars from different threads.
pub struct MultiProgress {
    objects: usize,
    draw_target: DrawTarget,
    tx: Sender<(usize, DrawState)>,
    rx: Receiver<(usize, DrawState)>,
}

impl MultiProgress {
    /// Creates a new multi progress object that draws to stdout.
    pub fn new() -> MultiProgress {
        let (tx, rx) = channel();
        MultiProgress {
            objects: 0,
            draw_target: DrawTarget::stdout(),
            tx: tx,
            rx: rx,
        }
    }

    /// Adds a progress bar.
    ///
    /// The progress bar added will have the draw target changed to a
    /// remote draw target that is intercepted by the multi progress
    /// object.
    pub fn add(&mut self, bar: ProgressBar) -> ProgressBar {
        bar.set_draw_target(DrawTarget::Remote(self.objects,
                                               self.tx.clone()));
        self.objects += 1;
        bar
    }

    /// Waits for all progress bars to report that they are finished.
    ///
    /// You need to call this as this will request the draw instructions
    /// from the remote progress bars.  Not calling this will deadlock
    /// your program.
    pub fn join(self) -> io::Result<()> {
        self.join_impl(false)
    }

    /// Works like `join` but clears the progress bar in the end.
    pub fn join_and_clear(self) -> io::Result<()> {
        self.join_impl(true)
    }

    fn join_impl(mut self, clear: bool) -> io::Result<()> {
        if self.objects == 0 {
            return Ok(());
        }

        let mut outstanding = repeat(true).take(self.objects as usize).collect::<Vec<_>>();
        let mut draw_states: Vec<Option<DrawState>> = outstanding.iter().map(|_| None).collect();

        while outstanding.iter().any(|&x| x) {
            let (idx, draw_state) = self.rx.recv().unwrap();
            let ts = draw_state.ts;

            if draw_state.finished {
                outstanding[idx] = false;
            }
            draw_states[idx] = Some(draw_state);

            let mut lines = vec![];
            for draw_state_opt in draw_states.iter() {
                if let Some(ref draw_state) = *draw_state_opt {
                    lines.extend_from_slice(&draw_state.lines[..]);
                }
            }

            self.draw_target.apply_draw_state(DrawState {
                lines: lines,
                finished: false,
                ts: ts,
            })?;
        }

        if clear {
            self.draw_target.apply_draw_state(DrawState {
                lines: vec![],
                finished: true,
                ts: Instant::now(),
            })?;
        }

        Ok(())
    }
}

impl Drop for ProgressBar {
    fn drop(&mut self) {
        if self.state.read().is_finished() {
            return;
        }
        self.update_state(|mut state| {
            state.status = Status::DoneHidden;
        });
    }
}
<commit_msg>Add a way to set the refresh rate<commit_after>use std::io;
use std::iter::repeat;
use std::borrow::Cow;
use std::cell::RefCell;
use std::time::{Duration, Instant};
use std::sync::mpsc::{channel, Sender, Receiver};

use parking_lot::RwLock;

use term::{Term, terminal_size};
use utils::expand_template;

/// Controls the rendering style of progress bars.
pub struct Style {
    pub tick_chars: Vec<char>,
    pub progress_chars: Vec<char>,
    pub bar_template: Cow<'static, str>,
    pub spinner_template: Cow<'static, str>,
}

/// The drawn state of an element.
#[derive(Clone)]
pub struct DrawState {
    /// The lines to print (can contain ANSI codes)
    pub lines: Vec<String>,
    /// True if the bar no longer needs drawing.
    pub finished: bool,
    /// Time when the draw state was created.
    pub ts: Instant,
}

enum Status {
    InProgress,
    DoneVisible,
    DoneHidden,
}

/// Target for draw operations
pub enum DrawTarget {
    /// Draws into a terminal
    Term(Term, Option<DrawState>, Option<Duration>),
    /// Draws to a remote receiver
    Remote(usize, Sender<(usize, DrawState)>),
    /// Do not draw at all
    Hidden,
}

impl DrawTarget {
    /// Draw to a terminal, optionally with a refresh rate.
    pub fn to_term(term: Term, refresh_rate: Option<u64>) -> DrawTarget {
        let rate = refresh_rate.map(|x| Duration::from_millis(1000 / x));
        DrawTarget::Term(term, None, rate)
    }

    /// Draw to a buffered stdout terminal at a max of 30 times a second.
    pub fn stdout() -> DrawTarget {
        DrawTarget::to_term(Term::buffered_stdout(), Some(30))
    }

    /// Draw to a buffered stderr terminal at a max of 30 times a second.
    pub fn stderr() -> DrawTarget {
        DrawTarget::to_term(Term::buffered_stderr(), Some(30))
    }

    /// Apply the given draw state (draws it).
    pub fn apply_draw_state(&mut self, draw_state: DrawState) -> io::Result<()> {
        match *self {
            DrawTarget::Term(ref term, ref mut last_state, rate) => {
                let last_draw = last_state.as_ref().map(|x| x.ts);
                if draw_state.finished ||
                   rate.is_none() ||
                   last_draw.is_none() ||
                   last_draw.unwrap().elapsed() > rate.unwrap() {
                    if let Some(ref last_state) = *last_state {
                        last_state.clear_term(term)?;
                    }
                    draw_state.draw_to_term(term)?;
                    term.flush()?;
                    *last_state = Some(draw_state);
                }
            }
            DrawTarget::Remote(idx, ref chan) => {
                chan.send((idx, draw_state)).unwrap();
            }
            DrawTarget::Hidden => {}
        }
        Ok(())
    }
}

impl DrawState {
    pub fn clear_term(&self, term: &Term) -> io::Result<()> {
        term.clear_last_lines(self.lines.len())
    }

    pub fn draw_to_term(&self, term: &Term) -> io::Result<()> {
        for line in &self.lines {
            term.write_line(line)?;
        }
        Ok(())
    }
}

impl Default for Style {
    fn default() -> Style {
        Style {
            tick_chars: "⠁⠁⠉⠙⠚⠒⠂⠂⠒⠲⠴⠤⠄⠄⠤⠠⠠⠤⠦⠖⠒⠐⠐⠒⠓⠋⠉⠈⠈ ".chars().collect(),
            progress_chars: "██░".chars().collect(),
            bar_template: Cow::Borrowed("{msg}\n{wide_bar} {pos}/{len}"),
            spinner_template: Cow::Borrowed("{spinner} {msg}"),
        }
    }
}

impl Style {

    /// Returns the tick char for a given number.
    pub fn get_tick_char(&self, idx: u64) -> char {
        self.tick_chars[(idx as usize) % (self.tick_chars.len() - 1)]
    }

    /// Returns the tick char for the finished state.
    pub fn get_final_tick_char(&self) -> char {
        self.tick_chars[self.tick_chars.len() - 1]
    }

    pub fn format_bar(&self, state: &ProgressState, width: usize) -> String {
        let pct = state.percent();
        let fill = (pct * width as f32) as usize;
        let bar = repeat(state.style.progress_chars[0]).take(fill).collect::<String>();
        let rest = repeat(state.style.progress_chars[2]).take(width - fill).collect::<String>();
        format!("{}{}", bar, rest)
    }

    pub fn format_state(&self, state: &ProgressState) -> Vec<String> {
        let (pos, len) = state.position();
        let mut rv = vec![];

        for line in self.bar_template.lines() {
            let need_wide_bar = RefCell::new(false);

            let s = expand_template(line, |key| {
                if key == "wide_bar" {
                    *need_wide_bar.borrow_mut() = true;
                    "\x00".into()
                } else if key == "bar" {
                    // XXX: width?
                    self.format_bar(state, 20)
                } else if key == "msg" {
                    state.message().to_string()
                } else if key == "pos" {
                    pos.to_string()
                } else if key == "len" {
                    len.to_string()
                } else {
                    "".into()
                }
            });

            rv.push(if *need_wide_bar.borrow() {
                let total_width = state.width();
                let bar_width = total_width - s.len() - 1;
                s.replace("\x00", &self.format_bar(state, bar_width))
            } else {
                s.to_string()
            });
        }

        rv
    }
}

/// The state of a progress bar at a moment in time.
pub struct ProgressState {
    style: Style,
    draw_target: DrawTarget,
    width: Option<u16>,
    message: String,
    pos: u64,
    len: u64,
    tick: u64,
    status: Status,
}

impl ProgressState {
    /// Returns the character that should be drawn for the
    /// current spinner character.
    pub fn current_tick_char(&self) -> char {
        if self.is_finished() {
            self.style.get_final_tick_char()
        } else {
            self.style.get_tick_char(self.tick)
        }
    }

    /// Indicates that a spinner should be drawn.
    pub fn has_spinner(&self) -> bool {
        self.tick != !0
    }

    /// Indicates that a progress bar should be drawn.
    pub fn has_progress(&self) -> bool {
        self.len != !0
    }

    /// Indicates that the progress bar finished.
    pub fn is_finished(&self) -> bool {
        match self.status {
            Status::InProgress => false,
            Status::DoneVisible => true,
            Status::DoneHidden => true,
        }
    }

    /// Returns `false` if the progress bar should no longer be
    /// drawn.
    pub fn should_render(&self) -> bool {
        match self.status {
            Status::DoneHidden => false,
            _ => true,
        }
    }

    /// Returns the completion in percent
    pub fn percent(&self) -> f32 {
        if self.len == !0 {
            0.0
        } else {
            self.pos as f32 / self.len as f32
        }
    }

    /// Returns the position of the status bar as `(pos, len)` tuple.
    pub fn position(&self) -> (u64, u64) {
        (self.pos, self.len)
    }

    /// Returns the current message of the progress bar.
    pub fn message(&self) -> &str {
        &self.message
    }

    /// The entire draw width
    pub fn width(&self) -> usize {
        if let Some(width) = self.width {
            width as usize
        } else {
            terminal_size().map(|x| x.1).unwrap_or(74) as usize
        }
    }
}

/// A progress bar or spinner.
pub struct ProgressBar {
    state: RwLock<ProgressState>,
}

impl ProgressBar {
    /// Creates a new progress bar with a given length.
    ///
    /// This progress bar by default draws directly to stdout.
    pub fn new(len: u64) -> ProgressBar {
        ProgressBar {
            state: RwLock::new(ProgressState {
                style: Default::default(),
                draw_target: DrawTarget::stdout(),
                width: None,
                message: "".into(),
                pos: 0,
                len: len,
                tick: !0,
                status: Status::InProgress,
            }),
        }
    }

    /// Creates a new spinner.
    ///
    /// This spinner by default draws directly to stdout.
    pub fn new_spinner() -> ProgressBar {
        let rv = ProgressBar::new(!0);
        rv.enable_spinner();
        rv
    }

    /// Enables the spinner.
    ///
    /// This is obviously enabled by default if you create a progress
    /// bar with `new_spinner` but optionally a spinner can be added
    /// to a progress bar itself.
    pub fn enable_spinner(&self) {
        self.update_state(|mut state| {
            if state.tick == !0 {
                state.tick = 0;
            }
        });
    }

    /// Disables a spinner.
    ///
    /// This should not be called if the progress bar is a spinner itself.
    pub fn disable_spinner(&self) {
        self.update_state(|mut state| {
            if state.tick != !0 {
                state.tick = !0;
            }
        });
    }

    /// Manually ticks the spinner or progress bar.
    ///
    /// This automatically happens on any other change to a progress bar.
    pub fn tick(&self) {
        self.update_state(|mut state| {
            if state.tick == !0 {
                state.tick = 0;
            } else {
                state.tick += 1;
            }
        });
    }

    /// Advances the position of a progress bar by delta.
    pub fn inc(&self, delta: u64) {
        self.update_state(|mut state| {
            state.pos += delta;
            if state.tick != !0 {
                state.tick += 1;
            }
        })
    }

    /// Sets the length of the progress bar.
    pub fn set_length(&self, len: u64) {
        self.update_state(|mut state| {
            state.len = len;
        })
    }

    /// Sets the current message of the progress bar.
    pub fn set_message(&self, msg: &str) {
        let msg = msg.to_string();
        self.update_state(|mut state| {
            state.message = msg;
        })
    }

    /// Finishes the progress bar and sets a message.
    pub fn finish_with_message(&self, msg: &str) {
        let msg = msg.to_string();
        self.update_state(|mut state| {
            state.message = msg;
            state.pos = state.len;
            state.status = Status::DoneVisible;
        });
    }

    /// Finishes the progress bar and completely clears it.
    pub fn finish_and_clear(&self) {
        self.update_state(|mut state| {
            state.pos = state.len;
            state.status = Status::DoneHidden;
        });
    }

    /// Sets a different draw target for the progress bar.
    ///
    /// This can be used to draw the progress bar to stderr
    /// for instance:
    ///
    /// ```rust,no_run
    /// # use indicatif::{ProgressBar, DrawTarget};
    /// let pb = ProgressBar::new(100);
    /// pb.set_draw_target(DrawTarget::stderr());
    /// ```
    pub fn set_draw_target(&self, target: DrawTarget) {
        self.state.write().draw_target = target;
    }

    fn update_state<F: FnOnce(&mut ProgressState)>(&self, f: F) {
        {
            let mut state = self.state.write();
            f(&mut state);
        }
        self.draw().ok();
    }

    fn draw(&self) -> io::Result<()> {
        let mut state = self.state.write();
        let draw_state = DrawState {
            lines: if state.should_render() {
                state.style.format_state(&*state)
            } else {
                vec![]
            },
            finished: state.is_finished(),
            ts: Instant::now(),
        };
        state.draw_target.apply_draw_state(draw_state)
    }
}

/// Manages multiple progress bars from different threads.
pub struct MultiProgress {
    objects: usize,
    draw_target: DrawTarget,
    tx: Sender<(usize, DrawState)>,
    rx: Receiver<(usize, DrawState)>,
}

impl MultiProgress {
    /// Creates a new multi progress object that draws to stdout.
    pub fn new() -> MultiProgress {
        let (tx, rx) = channel();
        MultiProgress {
            objects: 0,
            draw_target: DrawTarget::stdout(),
            tx: tx,
            rx: rx,
        }
    }

    pub fn set_draw_target(&mut self, target: DrawTarget) {
        self.draw_target = target;
    }

    /// Adds a progress bar.
    ///
    /// The progress bar added will have the draw target changed to a
    /// remote draw target that is intercepted by the multi progress
    /// object.
    pub fn add(&mut self, bar: ProgressBar) -> ProgressBar {
        bar.set_draw_target(DrawTarget::Remote(self.objects,
                                               self.tx.clone()));
        self.objects += 1;
        bar
    }

    /// Waits for all progress bars to report that they are finished.
    ///
    /// You need to call this as this will request the draw instructions
    /// from the remote progress bars.  Not calling this will deadlock
    /// your program.
    pub fn join(self) -> io::Result<()> {
        self.join_impl(false)
    }

    /// Works like `join` but clears the progress bar in the end.
    pub fn join_and_clear(self) -> io::Result<()> {
        self.join_impl(true)
    }

    fn join_impl(mut self, clear: bool) -> io::Result<()> {
        if self.objects == 0 {
            return Ok(());
        }

        let mut outstanding = repeat(true).take(self.objects as usize).collect::<Vec<_>>();
        let mut draw_states: Vec<Option<DrawState>> = outstanding.iter().map(|_| None).collect();

        while outstanding.iter().any(|&x| x) {
            let (idx, draw_state) = self.rx.recv().unwrap();
            let ts = draw_state.ts;

            if draw_state.finished {
                outstanding[idx] = false;
            }
            draw_states[idx] = Some(draw_state);

            let mut lines = vec![];
            for draw_state_opt in draw_states.iter() {
                if let Some(ref draw_state) = *draw_state_opt {
                    lines.extend_from_slice(&draw_state.lines[..]);
                }
            }

            self.draw_target.apply_draw_state(DrawState {
                lines: lines,
                finished: false,
                ts: ts,
            })?;
        }

        if clear {
            self.draw_target.apply_draw_state(DrawState {
                lines: vec![],
                finished: true,
                ts: Instant::now(),
            })?;
        }

        Ok(())
    }
}

impl Drop for ProgressBar {
    fn drop(&mut self) {
        if self.state.read().is_finished() {
            return;
        }
        self.update_state(|mut state| {
            state.status = Status::DoneHidden;
        });
    }
}
<|endoftext|>