<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/.

extern crate libc;

use std::collections::HashMap;
use std::os::unix::io::{AsRawFd, RawFd};
use std::path::{Path, PathBuf};

use devicemapper::{DM, Device, DmNameBuf};

use super::super::engine::{Engine, Eventable, Pool};
use super::super::errors::{EngineError, EngineResult, ErrorEnum};
use super::super::structures::Table;
use super::super::types::{DevUuid, Name, PoolUuid, Redundancy, RenameAction};

use super::backstore::{find_all, is_stratis_device};
use super::cleanup::teardown_pools;
use super::devlinks;
use super::pool::StratPool;

pub const DEV_PATH: &str = "/dev/stratis";
const REQUIRED_DM_MINOR_VERSION: u32 = 37;

#[derive(Debug, PartialEq, Eq)]
pub enum DevOwnership {
    Ours(PoolUuid, DevUuid),
    Unowned,
    Theirs,
}

#[derive(Debug)]
pub struct StratEngine {
    pools: Table<StratPool>,

    // Map of stratis devices that have been found but one or more stratis block devices are missing
    // which prevents the associated pools from being setup.
    incomplete_pools: HashMap<PoolUuid, HashMap<Device, PathBuf>>,

    // Maps name of DM devices we are watching to the most recent event number
    // we've handled for each
    watched_dev_last_event_nrs: HashMap<DmNameBuf, u32>,
}

impl Eventable for DM {
    /// Get file we'd like to have monitored for activity
    fn get_pollable_fd(&mut self) -> RawFd {
        self.file().as_raw_fd()
    }

    fn clear_event(&mut self) -> EngineResult<()> {
        self.arm_poll()?;

        Ok(())
    }
}

impl StratEngine {
    /// Setup a StratEngine.
    /// 1. Verify the existence of Stratis /dev directory.
    /// 2. Setup all the pools belonging to the engine.
    ///    a. Places any devices which belong to a pool, but are not complete
    ///       in the incomplete pools data structure.
    ///
    /// Returns an error if the kernel doesn't support required DM features.
    /// Returns an error if there was an error reading device nodes.
    /// Returns an error and tears down all the pools if we find a pool with a
    /// duplicate name found.
    pub fn initialize() -> EngineResult<StratEngine> {
        let dm = DM::new()?;
        let minor_dm_version = dm.version()?.1;
        if minor_dm_version < REQUIRED_DM_MINOR_VERSION {
            let err_msg = format!("Requires DM minor version {} but kernel only supports {}",
                                  REQUIRED_DM_MINOR_VERSION,
                                  minor_dm_version);
            return Err(EngineError::Engine(ErrorEnum::Error, err_msg));
        }

        devlinks::setup_dev_path()?;

        let pools = find_all()?;

        let mut table = Table::default();
        let mut incomplete_pools = HashMap::new();

        for (pool_uuid, devices) in &pools {
            match StratPool::setup(&dm, *pool_uuid, devices) {
                Ok((pool_name, pool)) => {
                    let evicted = table.insert(pool_name, *pool_uuid, pool);
                    if evicted.is_some() {
                        // TODO: update state machine on failure.
                        let _ = teardown_pools(&dm, table);
                        let err_msg = "found two pools with the same name";
                        return Err(EngineError::Engine(ErrorEnum::Invalid, err_msg.into()));
                    }
                }
                Err(_) => {
                    // Unable to setup pool, it might become complete later.
                    incomplete_pools.insert(*pool_uuid, devices.clone());
                }
            }
        }

        let engine = StratEngine {
            pools: table,
            incomplete_pools,
            watched_dev_last_event_nrs: HashMap::new(),
        };

        devlinks::setup_devlinks(engine.pools().iter())?;

        Ok(engine)
    }

    /// Teardown Stratis, preparatory to a shutdown.
    #[cfg(test)]
    pub fn teardown(self) -> EngineResult<()> {
        teardown_pools(&DM::new()?, self.pools)
    }
}

impl Engine for StratEngine {
    fn create_pool(&mut self,
                   name: &str,
                   blockdev_paths: &[&Path],
                   redundancy: Option<u16>,
                   force: bool)
                   -> EngineResult<PoolUuid> {

        let redundancy = calculate_redundancy!(redundancy);

        if self.pools.contains_name(name) {
            return Err(EngineError::Engine(ErrorEnum::AlreadyExists, name.into()));
        }

        let dm = DM::new()?;
        let (uuid, pool) = StratPool::initialize(&dm, name, blockdev_paths, redundancy, force)?;

        let name = Name::new(name.to_owned());
        devlinks::pool_added(&name)?;
        self.pools.insert(name, uuid, pool);
        Ok(uuid)
    }

    /// Evaluate a device node & devicemapper::Device to see if it's a valid
    /// stratis device.  If all the devices are present in the pool and the pool isn't already
    /// up and running, it will get setup and the pool uuid will be returned.
    ///
    /// Returns an error if the status of the block device can not be evaluated.
    fn block_evaluate(&mut self,
                      dev_node: PathBuf,
                      device: Device)
                      -> EngineResult<Option<PoolUuid>> {
        let dm = DM::new()?;
        if let Some(pool_uuid) = is_stratis_device(&dev_node)? {
            if self.pools.contains_uuid(pool_uuid) {
                // TODO: Handle the case where we have found a device for an already active pool
                // ref. https://github.com/stratis-storage/stratisd/issues/748
                warn!("udev add: pool {} is already known, ignoring device {:?}!",
                      pool_uuid,
                      dev_node);
                Ok(None)
            } else {

                let mut devices = self.incomplete_pools
                    .remove(&pool_uuid)
                    .or_else(|| Some(HashMap::new()))
                    .expect("We just retrieved or created a HashMap");
                devices.insert(device, dev_node);

                let rc = if let Ok((pool_name, pool)) =
                    StratPool::setup(&dm, pool_uuid, &devices) {
                    // We know we have a unique uuid, but the pools table requires a unique name too
                    // so we will ensure unique before we insert as we don't want to change the
                    // existing state if we have a conflict.
                    //
                    // We will also _not_ exit the daemon as we do in initialize as we were
                    // previously up and running for some duration of time.
                    if !self.pools.contains_name(&pool_name) {
                        self.pools.insert(pool_name, pool_uuid, pool);
                        Some(pool_uuid)
                    } else {
                        let dev_paths = devices
                            .values()
                            .map(|p| p.to_str().expect("Expecting valid device path!"))
                            .collect::<Vec<&str>>()
                            .join(", ");

                        self.incomplete_pools.insert(pool_uuid, devices);

                        error!("udev add: duplicate pool name {:?} for uuid {:?}, \
                                devices[{}], failing to setup complete pool!",
                               pool_name,
                               pool_uuid,
                               dev_paths);
                        if let Err(e) = pool.teardown(&dm) {
                            error!("Error while tearing down pool with duplicate name! {:?}!",
                                   e);
                        }
                        None
                    }
                } else {
                    self.incomplete_pools.insert(pool_uuid, devices);
                    None
                };
                Ok(rc)
            }
        } else {
            Ok(None)
        }
    }

    fn destroy_pool(&mut self, uuid: PoolUuid) -> EngineResult<bool> {
        if let Some((_, pool)) = self.pools.get_by_uuid(uuid) {
            if pool.has_filesystems() {
                return Err(EngineError::Engine(ErrorEnum::Busy,
                                               "filesystems remaining on pool".into()));
            };
        } else {
            return Ok(false);
        }

        let (pool_name, pool) =
            self.pools
                .remove_by_uuid(uuid)
                .expect("Must succeed since self.pools.get_by_uuid() returned a value");

        pool.destroy()?;
        devlinks::pool_removed(&pool_name)?;
        Ok(true)
    }

    fn rename_pool(&mut self, uuid: PoolUuid, new_name: &str) -> EngineResult<RenameAction> {
        let old_name = rename_pool_pre!(self; uuid; new_name);

        let (_, mut pool) =
            self.pools
                .remove_by_uuid(uuid)
                .expect("Must succeed since self.pools.get_by_uuid() returned a value");

        let new_name = Name::new(new_name.to_owned());
        if let Err(err) = pool.write_metadata(&new_name) {
            self.pools.insert(old_name, uuid, pool);
            Err(err)
        } else {
            self.pools.insert(new_name.clone(), uuid, pool);
            devlinks::pool_renamed(&old_name, &new_name)?;
            Ok(RenameAction::Renamed)
        }
    }

    fn get_pool(&self, uuid: PoolUuid) -> Option<(Name, &Pool)> {
        get_pool!(self; uuid)
    }

    fn get_mut_pool(&mut self, uuid: PoolUuid) -> Option<(Name, &mut Pool)> {
        get_mut_pool!(self; uuid)
    }

    fn configure_simulator(&mut self, _denominator: u32) -> EngineResult<()> {
        Ok(()) // we're not the simulator and not configurable, so just say ok
    }

    fn check(&mut self) -> () {
        check_engine!(self);
    }

    fn pools(&self) -> Vec<(Name, PoolUuid, &Pool)> {
        self.pools
            .iter()
            .map(|(name, uuid, pool)| (name.clone(), *uuid, pool as &Pool))
            .collect()
    }

    fn get_eventable(&mut self) -> EngineResult<Option<Box<Eventable>>> {
        Ok(Some(Box::new(DM::new()?)))
    }

    fn evented(&mut self) -> EngineResult<()> {
        let device_list: HashMap<_, _> = DM::new()?
            .list_devices()?
            .into_iter()
            .map(|(dm_name, _, event_nr)| {
                     (dm_name, event_nr.expect("Supported DM versions always provide a value"))
                 })
            .collect();

        for (pool_name, _, pool) in &mut self.pools {
            for dm_name in pool.get_eventing_dev_names() {
                if device_list.get(&dm_name) > self.watched_dev_last_event_nrs.get(&dm_name) {
                    pool.event_on(pool_name, &dm_name)?;
                }
            }
        }

        self.watched_dev_last_event_nrs = device_list;

        Ok(())
    }
}

#[cfg(test)]
mod test {
    use std::fs::remove_dir_all;

    use super::super::tests::{loopbacked, real};

    use super::*;

    /// Verify that a pool rename causes the pool metadata to get the new name.
    fn test_pool_rename(paths: &[&Path]) {
        let mut engine = StratEngine::initialize().unwrap();

        let name1 = "name1";
        let uuid1 = engine.create_pool(&name1, paths, None, false).unwrap();

        let name2 = "name2";
        let action = engine.rename_pool(uuid1, name2).unwrap();

        assert_eq!(action, RenameAction::Renamed);
        engine.teardown().unwrap();

        let engine = StratEngine::initialize().unwrap();
        let pool_name: String = engine.get_pool(uuid1).unwrap().0.to_owned();
        assert_eq!(pool_name, name2);
    }

    #[test]
    pub fn loop_test_pool_rename() {
        loopbacked::test_with_spec(loopbacked::DeviceLimits::Range(1, 3), test_pool_rename);
    }

    #[test]
    pub fn real_test_pool_rename() {
        real::test_with_spec(real::DeviceLimits::AtLeast(1), test_pool_rename);
    }

    /// Test engine setup.
    /// 1. Create two pools.
    /// 2. Verify that both exist.
    /// 3. Teardown the engine.
    /// 4. Initialize the engine.
    /// 5. Verify that pools can be found again.
    /// 6. Teardown the engine and remove "/dev/stratis".
    /// 7. Initialize the engine one more time.
    /// 8. Verify that both pools are found and that there are no incomplete pools.
    fn test_setup(paths: &[&Path]) {
        assert!(paths.len() > 1);

        let (paths1, paths2) = paths.split_at(paths.len() / 2);

        let mut engine = StratEngine::initialize().unwrap();

        let name1 = "name1";
        let uuid1 = engine.create_pool(&name1, paths1, None, false).unwrap();

        let name2 = "name2";
        let uuid2 = engine.create_pool(&name2, paths2, None, false).unwrap();

        assert!(engine.get_pool(uuid1).is_some());
        assert!(engine.get_pool(uuid2).is_some());

        engine.teardown().unwrap();

        let engine = StratEngine::initialize().unwrap();

        assert!(engine.get_pool(uuid1).is_some());
        assert!(engine.get_pool(uuid2).is_some());

        engine.teardown().unwrap();
        remove_dir_all(DEV_PATH).unwrap();

        let engine = StratEngine::initialize().unwrap();
        assert_eq!(engine.incomplete_pools, HashMap::new());

        assert!(engine.get_pool(uuid1).is_some());
        assert!(engine.get_pool(uuid2).is_some());

        engine.teardown().unwrap();
    }

    #[test]
    pub fn loop_test_setup() {
        loopbacked::test_with_spec(loopbacked::DeviceLimits::Range(2, 3), test_setup);
    }

    #[test]
    pub fn real_test_setup() {
        real::test_with_spec(real::DeviceLimits::AtLeast(2), test_setup);
    }
}
<commit_msg>strat_engine/engine.rs: Add warn! on pool setup fail<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/.

extern crate libc;

use std::collections::HashMap;
use std::os::unix::io::{AsRawFd, RawFd};
use std::path::{Path, PathBuf};

use devicemapper::{DM, Device, DmNameBuf};

use super::super::engine::{Engine, Eventable, Pool};
use super::super::errors::{EngineError, EngineResult, ErrorEnum};
use super::super::structures::Table;
use super::super::types::{DevUuid, Name, PoolUuid, Redundancy, RenameAction};

use super::backstore::{find_all, is_stratis_device};
use super::cleanup::teardown_pools;
use super::devlinks;
use super::pool::StratPool;

pub const DEV_PATH: &str = "/dev/stratis";
const REQUIRED_DM_MINOR_VERSION: u32 = 37;

#[derive(Debug, PartialEq, Eq)]
pub enum DevOwnership {
    Ours(PoolUuid, DevUuid),
    Unowned,
    Theirs,
}

#[derive(Debug)]
pub struct StratEngine {
    pools: Table<StratPool>,

    // Map of stratis devices that have been found but one or more stratis block devices are missing
    // which prevents the associated pools from being setup.
    incomplete_pools: HashMap<PoolUuid, HashMap<Device, PathBuf>>,

    // Maps name of DM devices we are watching to the most recent event number
    // we've handled for each
    watched_dev_last_event_nrs: HashMap<DmNameBuf, u32>,
}

impl Eventable for DM {
    /// Get file we'd like to have monitored for activity
    fn get_pollable_fd(&mut self) -> RawFd {
        self.file().as_raw_fd()
    }

    fn clear_event(&mut self) -> EngineResult<()> {
        self.arm_poll()?;

        Ok(())
    }
}

impl StratEngine {
    /// Setup a StratEngine.
    /// 1. Verify the existence of Stratis /dev directory.
    /// 2. Setup all the pools belonging to the engine.
    ///    a. Places any devices which belong to a pool, but are not complete
    ///       in the incomplete pools data structure.
    ///
    /// Returns an error if the kernel doesn't support required DM features.
    /// Returns an error if there was an error reading device nodes.
    /// Returns an error and tears down all the pools if we find a pool with a
    /// duplicate name found.
    pub fn initialize() -> EngineResult<StratEngine> {
        let dm = DM::new()?;
        let minor_dm_version = dm.version()?.1;
        if minor_dm_version < REQUIRED_DM_MINOR_VERSION {
            let err_msg = format!("Requires DM minor version {} but kernel only supports {}",
                                  REQUIRED_DM_MINOR_VERSION,
                                  minor_dm_version);
            return Err(EngineError::Engine(ErrorEnum::Error, err_msg));
        }

        devlinks::setup_dev_path()?;

        let pools = find_all()?;

        let mut table = Table::default();
        let mut incomplete_pools = HashMap::new();

        for (pool_uuid, devices) in &pools {
            match StratPool::setup(&dm, *pool_uuid, devices) {
                Ok((pool_name, pool)) => {
                    let evicted = table.insert(pool_name, *pool_uuid, pool);
                    if evicted.is_some() {
                        // TODO: update state machine on failure.
                        let _ = teardown_pools(&dm, table);
                        let err_msg = "found two pools with the same name";
                        return Err(EngineError::Engine(ErrorEnum::Invalid, err_msg.into()));
                    }
                }
                Err(e) => {
                    // TODO: Add ability to distinguish expected vs. unexpected errors,
                    // expected error would be caused by missing block device(s).
                    warn!("Pool {} failed to setup, reason {:?}.", pool_uuid, e);
                    // Unable to setup pool, it might become complete later.
                    incomplete_pools.insert(*pool_uuid, devices.clone());
                }
            }
        }

        let engine = StratEngine {
            pools: table,
            incomplete_pools,
            watched_dev_last_event_nrs: HashMap::new(),
        };

        devlinks::setup_devlinks(engine.pools().iter())?;

        Ok(engine)
    }

    /// Teardown Stratis, preparatory to a shutdown.
    #[cfg(test)]
    pub fn teardown(self) -> EngineResult<()> {
        teardown_pools(&DM::new()?, self.pools)
    }
}

impl Engine for StratEngine {
    fn create_pool(&mut self,
                   name: &str,
                   blockdev_paths: &[&Path],
                   redundancy: Option<u16>,
                   force: bool)
                   -> EngineResult<PoolUuid> {

        let redundancy = calculate_redundancy!(redundancy);

        if self.pools.contains_name(name) {
            return Err(EngineError::Engine(ErrorEnum::AlreadyExists, name.into()));
        }

        let dm = DM::new()?;
        let (uuid, pool) = StratPool::initialize(&dm, name, blockdev_paths, redundancy, force)?;

        let name = Name::new(name.to_owned());
        devlinks::pool_added(&name)?;
        self.pools.insert(name, uuid, pool);
        Ok(uuid)
    }

    /// Evaluate a device node & devicemapper::Device to see if it's a valid
    /// stratis device.  If all the devices are present in the pool and the pool isn't already
    /// up and running, it will get setup and the pool uuid will be returned.
    ///
    /// Returns an error if the status of the block device can not be evaluated.
    fn block_evaluate(&mut self,
                      dev_node: PathBuf,
                      device: Device)
                      -> EngineResult<Option<PoolUuid>> {
        let dm = DM::new()?;
        if let Some(pool_uuid) = is_stratis_device(&dev_node)? {
            if self.pools.contains_uuid(pool_uuid) {
                // TODO: Handle the case where we have found a device for an already active pool
                // ref. https://github.com/stratis-storage/stratisd/issues/748
                warn!("udev add: pool {} is already known, ignoring device {:?}!",
                      pool_uuid,
                      dev_node);
                Ok(None)
            } else {

                let mut devices = self.incomplete_pools
                    .remove(&pool_uuid)
                    .or_else(|| Some(HashMap::new()))
                    .expect("We just retrieved or created a HashMap");
                devices.insert(device, dev_node);

                let rc = match StratPool::setup(&dm, pool_uuid, &devices) {
                    Ok((pool_name, pool)) => {
                        // We know we have a unique uuid, but the pools table requires a unique
                        // name too so we will ensure unique before we insert as we don't want to
                        // change the existing state if we have a conflict.
                        //
                        // We will also _not_ exit the daemon as we do in initialize as we were
                        // previously up and running for some duration of time.
                        if !self.pools.contains_name(&pool_name) {
                            self.pools.insert(pool_name, pool_uuid, pool);
                            Some(pool_uuid)
                        } else {
                            let dev_paths = devices
                                .values()
                                .map(|p| p.to_str().expect("Expecting valid device path!"))
                                .collect::<Vec<&str>>()
                                .join(", ");

                            self.incomplete_pools.insert(pool_uuid, devices);

                            error!("udev add: duplicate pool name {:?} for uuid {:?}, \
                                devices[{}], failing to setup complete pool.",
                                   pool_name,
                                   pool_uuid,
                                   dev_paths);
                            if let Err(e) = pool.teardown(&dm) {
                                error!("Error while tearing down pool with duplicate name {:?}.",
                                       e);
                            }
                            None
                        }
                    }
                    Err(e) => {
                        warn!("udev add: pool {} failed to setup, reason {:?}.",
                              pool_uuid,
                              e);
                        self.incomplete_pools.insert(pool_uuid, devices);
                        None
                    }
                };
                Ok(rc)
            }
        } else {
            Ok(None)
        }
    }

    fn destroy_pool(&mut self, uuid: PoolUuid) -> EngineResult<bool> {
        if let Some((_, pool)) = self.pools.get_by_uuid(uuid) {
            if pool.has_filesystems() {
                return Err(EngineError::Engine(ErrorEnum::Busy,
                                               "filesystems remaining on pool".into()));
            };
        } else {
            return Ok(false);
        }

        let (pool_name, pool) =
            self.pools
                .remove_by_uuid(uuid)
                .expect("Must succeed since self.pools.get_by_uuid() returned a value");

        pool.destroy()?;
        devlinks::pool_removed(&pool_name)?;
        Ok(true)
    }

    fn rename_pool(&mut self, uuid: PoolUuid, new_name: &str) -> EngineResult<RenameAction> {
        let old_name = rename_pool_pre!(self; uuid; new_name);

        let (_, mut pool) =
            self.pools
                .remove_by_uuid(uuid)
                .expect("Must succeed since self.pools.get_by_uuid() returned a value");

        let new_name = Name::new(new_name.to_owned());
        if let Err(err) = pool.write_metadata(&new_name) {
            self.pools.insert(old_name, uuid, pool);
            Err(err)
        } else {
            self.pools.insert(new_name.clone(), uuid, pool);
            devlinks::pool_renamed(&old_name, &new_name)?;
            Ok(RenameAction::Renamed)
        }
    }

    fn get_pool(&self, uuid: PoolUuid) -> Option<(Name, &Pool)> {
        get_pool!(self; uuid)
    }

    fn get_mut_pool(&mut self, uuid: PoolUuid) -> Option<(Name, &mut Pool)> {
        get_mut_pool!(self; uuid)
    }

    fn configure_simulator(&mut self, _denominator: u32) -> EngineResult<()> {
        Ok(()) // we're not the simulator and not configurable, so just say ok
    }

    fn check(&mut self) -> () {
        check_engine!(self);
    }

    fn pools(&self) -> Vec<(Name, PoolUuid, &Pool)> {
        self.pools
            .iter()
            .map(|(name, uuid, pool)| (name.clone(), *uuid, pool as &Pool))
            .collect()
    }

    fn get_eventable(&mut self) -> EngineResult<Option<Box<Eventable>>> {
        Ok(Some(Box::new(DM::new()?)))
    }

    fn evented(&mut self) -> EngineResult<()> {
        let device_list: HashMap<_, _> = DM::new()?
            .list_devices()?
            .into_iter()
            .map(|(dm_name, _, event_nr)| {
                     (dm_name, event_nr.expect("Supported DM versions always provide a value"))
                 })
            .collect();

        for (pool_name, _, pool) in &mut self.pools {
            for dm_name in pool.get_eventing_dev_names() {
                if device_list.get(&dm_name) > self.watched_dev_last_event_nrs.get(&dm_name) {
                    pool.event_on(pool_name, &dm_name)?;
                }
            }
        }

        self.watched_dev_last_event_nrs = device_list;

        Ok(())
    }
}

#[cfg(test)]
mod test {
    use std::fs::remove_dir_all;

    use super::super::tests::{loopbacked, real};

    use super::*;

    /// Verify that a pool rename causes the pool metadata to get the new name.
    fn test_pool_rename(paths: &[&Path]) {
        let mut engine = StratEngine::initialize().unwrap();

        let name1 = "name1";
        let uuid1 = engine.create_pool(&name1, paths, None, false).unwrap();

        let name2 = "name2";
        let action = engine.rename_pool(uuid1, name2).unwrap();

        assert_eq!(action, RenameAction::Renamed);
        engine.teardown().unwrap();

        let engine = StratEngine::initialize().unwrap();
        let pool_name: String = engine.get_pool(uuid1).unwrap().0.to_owned();
        assert_eq!(pool_name, name2);
    }

    #[test]
    pub fn loop_test_pool_rename() {
        loopbacked::test_with_spec(loopbacked::DeviceLimits::Range(1, 3), test_pool_rename);
    }

    #[test]
    pub fn real_test_pool_rename() {
        real::test_with_spec(real::DeviceLimits::AtLeast(1), test_pool_rename);
    }

    /// Test engine setup.
    /// 1. Create two pools.
    /// 2. Verify that both exist.
    /// 3. Teardown the engine.
    /// 4. Initialize the engine.
    /// 5. Verify that pools can be found again.
    /// 6. Teardown the engine and remove "/dev/stratis".
    /// 7. Initialize the engine one more time.
    /// 8. Verify that both pools are found and that there are no incomplete pools.
    fn test_setup(paths: &[&Path]) {
        assert!(paths.len() > 1);

        let (paths1, paths2) = paths.split_at(paths.len() / 2);

        let mut engine = StratEngine::initialize().unwrap();

        let name1 = "name1";
        let uuid1 = engine.create_pool(&name1, paths1, None, false).unwrap();

        let name2 = "name2";
        let uuid2 = engine.create_pool(&name2, paths2, None, false).unwrap();

        assert!(engine.get_pool(uuid1).is_some());
        assert!(engine.get_pool(uuid2).is_some());

        engine.teardown().unwrap();

        let engine = StratEngine::initialize().unwrap();

        assert!(engine.get_pool(uuid1).is_some());
        assert!(engine.get_pool(uuid2).is_some());

        engine.teardown().unwrap();
        remove_dir_all(DEV_PATH).unwrap();

        let engine = StratEngine::initialize().unwrap();
        assert_eq!(engine.incomplete_pools, HashMap::new());

        assert!(engine.get_pool(uuid1).is_some());
        assert!(engine.get_pool(uuid2).is_some());

        engine.teardown().unwrap();
    }

    #[test]
    pub fn loop_test_setup() {
        loopbacked::test_with_spec(loopbacked::DeviceLimits::Range(2, 3), test_setup);
    }

    #[test]
    pub fn real_test_setup() {
        real::test_with_spec(real::DeviceLimits::AtLeast(2), test_setup);
    }
}
<|endoftext|>


<commit_before><commit_msg>Qualify ast types<commit_after><|endoftext|>


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

use target::TargetOptions;
use std::default::Default;

pub fn opts() -> TargetOptions {
    TargetOptions {
        linker: "cc".to_string(),
        dynamic_linking: true,
        executables: true,
        linker_is_gnu: true,
        has_rpath: true,
        pre_link_args: vec!(
            "-L/usr/local/lib".to_string(),
            "-L/usr/lib/gcc47".to_string(),
            // GNU-style linkers will use this to omit linking to libraries
            // which don't actually fulfill any relocations, but only for
            // libraries which follow this flag.  Thus, use it before
            // specifying libraries to link to.
            "-Wl,--as-needed".to_string(),
        ),
        position_independent_executables: true,
        archive_format: "bsd".to_string(),
        exe_allocation_crate: super::best_allocator(),
        .. Default::default()
    }
}
<commit_msg>fixes #27124 for DragonFly<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use target::TargetOptions;
use std::default::Default;

pub fn opts() -> TargetOptions {
    TargetOptions {
        linker: "cc".to_string(),
        dynamic_linking: true,
        executables: true,
        linker_is_gnu: true,
        has_rpath: true,
        pre_link_args: vec!(
            "-L/usr/local/lib".to_string(),
            "-L/usr/lib/gcc47".to_string(),
            // GNU-style linkers will use this to omit linking to libraries
            // which don't actually fulfill any relocations, but only for
            // libraries which follow this flag.  Thus, use it before
            // specifying libraries to link to.
            "-Wl,--as-needed".to_string(),
        ),
        position_independent_executables: true,
        archive_format: "gnu".to_string(),
        exe_allocation_crate: super::best_allocator(),
        .. Default::default()
    }
}
<|endoftext|>


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

use bit_set::BitSet;
use indexed_vec::Idx;
use std::collections::VecDeque;

/// A work queue is a handy data structure for tracking work left to
/// do. (For example, basic blocks left to process.) It is basically a
/// de-duplicating queue; so attempting to insert X if X is already
/// enqueued has no effect. This implementation assumes that the
/// elements are dense indices, so it can allocate the queue to size
/// and also use a bit set to track occupancy.
pub struct WorkQueue<T: Idx> {
    deque: VecDeque<T>,
    set: BitSet<T>,
}

impl<T: Idx> WorkQueue<T> {
    /// Create a new work queue with all the elements from (0..len).
    #[inline]
    pub fn with_all(len: usize) -> Self {
        WorkQueue {
            deque: (0..len).map(T::new).collect(),
            set: BitSet::new_filled(len),
        }
    }

    /// Create a new work queue that starts empty, where elements range from (0..len).
    #[inline]
    pub fn with_none(len: usize) -> Self {
        WorkQueue {
            deque: VecDeque::with_capacity(len),
            set: BitSet::new_empty(len),
        }
    }

    /// Attempt to enqueue `element` in the work queue. Returns false if it was already present.
    #[inline]
    pub fn insert(&mut self, element: T) -> bool {
        if self.set.insert(element) {
            self.deque.push_back(element);
            true
        } else {
            false
        }
    }

    /// Attempt to enqueue `element` in the work queue. Returns false if it was already present.
    #[inline]
    pub fn pop(&mut self) -> Option<T> {
        if let Some(element) = self.deque.pop_front() {
            self.set.remove(element);
            Some(element)
        } else {
            None
        }
    }

    /// True if nothing is enqueued.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.deque.is_empty()
    }
}
<commit_msg>Correct doc for WorkQueue<T>::pop().<commit_after>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use bit_set::BitSet;
use indexed_vec::Idx;
use std::collections::VecDeque;

/// A work queue is a handy data structure for tracking work left to
/// do. (For example, basic blocks left to process.) It is basically a
/// de-duplicating queue; so attempting to insert X if X is already
/// enqueued has no effect. This implementation assumes that the
/// elements are dense indices, so it can allocate the queue to size
/// and also use a bit set to track occupancy.
pub struct WorkQueue<T: Idx> {
    deque: VecDeque<T>,
    set: BitSet<T>,
}

impl<T: Idx> WorkQueue<T> {
    /// Create a new work queue with all the elements from (0..len).
    #[inline]
    pub fn with_all(len: usize) -> Self {
        WorkQueue {
            deque: (0..len).map(T::new).collect(),
            set: BitSet::new_filled(len),
        }
    }

    /// Create a new work queue that starts empty, where elements range from (0..len).
    #[inline]
    pub fn with_none(len: usize) -> Self {
        WorkQueue {
            deque: VecDeque::with_capacity(len),
            set: BitSet::new_empty(len),
        }
    }

    /// Attempt to enqueue `element` in the work queue. Returns false if it was already present.
    #[inline]
    pub fn insert(&mut self, element: T) -> bool {
        if self.set.insert(element) {
            self.deque.push_back(element);
            true
        } else {
            false
        }
    }

    /// Attempt to pop an element from the work queue.
    #[inline]
    pub fn pop(&mut self) -> Option<T> {
        if let Some(element) = self.deque.pop_front() {
            self.set.remove(element);
            Some(element)
        } else {
            None
        }
    }

    /// True if nothing is enqueued.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.deque.is_empty()
    }
}
<|endoftext|>


<commit_before><commit_msg>Assert formatted float parses back to the same value<commit_after><|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/. */

use dom::bindings::codegen::Bindings::RangeBinding;
use dom::bindings::codegen::Bindings::RangeBinding::RangeMethods;
use dom::bindings::error::Fallible;
use dom::bindings::global::GlobalRef;
use dom::bindings::js::{JSRef, Temporary};
use dom::bindings::utils::{Reflectable, Reflector, reflect_dom_object};

#[deriving(Encodable)]
pub struct Range {
    reflector_: Reflector
}

impl Range {
    pub fn new_inherited() -> Range {
        Range {
            reflector_: Reflector::new()
        }
    }

    pub fn new(global: &GlobalRef) -> Temporary<Range> {
        reflect_dom_object(box Range::new_inherited(), global, RangeBinding::Wrap)
    }

    pub fn Constructor(global: &GlobalRef) -> Fallible<Temporary<Range>> {
        Ok(Range::new(global))
    }
}

impl<'a> RangeMethods for JSRef<'a, Range> {
}

impl Reflectable for Range {
    fn reflector<'a>(&'a self) -> &'a Reflector {
        &self.reflector_
    }
}
<commit_msg>Pass a Document to Range::new.<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/. */

use dom::bindings::codegen::Bindings::RangeBinding;
use dom::bindings::codegen::Bindings::RangeBinding::RangeMethods;
use dom::bindings::codegen::Bindings::WindowBinding::WindowMethods;
use dom::bindings::error::Fallible;
use dom::bindings::global::{GlobalRef, Window};
use dom::bindings::js::{JSRef, Temporary};
use dom::bindings::utils::{Reflectable, Reflector, reflect_dom_object};
use dom::document::Document;

#[deriving(Encodable)]
pub struct Range {
    reflector_: Reflector
}

impl Range {
    pub fn new_inherited() -> Range {
        Range {
            reflector_: Reflector::new()
        }
    }

    pub fn new(document: &JSRef<Document>) -> Temporary<Range> {
        let window = document.window.root();
        reflect_dom_object(box Range::new_inherited(),
                           &Window(*window),
                           RangeBinding::Wrap)
    }

    pub fn Constructor(global: &GlobalRef) -> Fallible<Temporary<Range>> {
        let document = global.as_window().Document().root();
        Ok(Range::new(&*document))
    }
}

impl<'a> RangeMethods for JSRef<'a, Range> {
}

impl Reflectable for Range {
    fn reflector<'a>(&'a self) -> &'a Reflector {
        &self.reflector_
    }
}
<|endoftext|>


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

//! Compiles the profiler part of the `compiler-rt` library.
//!
//! See the build.rs for libcompiler_builtins crate for details.

extern crate gcc;

use std::env;
use std::path::Path;

fn main() {
    let target = env::var("TARGET").expect("TARGET was not set");
    let cfg = &mut gcc::Config::new();

    if target.contains("msvc") {
        // Don't pull in extra libraries on MSVC
        cfg.flag("/Zl");
    } else {
        // Turn off various features of gcc and such, mostly copying
        // compiler-rt's build system already
        cfg.flag("-fno-builtin");
        cfg.flag("-fvisibility=hidden");
        cfg.flag("-fomit-frame-pointer");
        cfg.flag("-ffreestanding");
        cfg.define("VISIBILITY_HIDDEN", None);
    }

    let profile_sources = &["GCDAProfiling.c",
                            "InstrProfiling.c",
                            "InstrProfilingBuffer.c",
                            "InstrProfilingFile.c",
                            "InstrProfilingMerge.c",
                            "InstrProfilingMergeFile.c",
                            "InstrProfilingPlatformDarwin.c",
                            "InstrProfilingPlatformLinux.c",
                            "InstrProfilingPlatformOther.c",
                            "InstrProfilingRuntime.cc",
                            "InstrProfilingUtil.c",
                            "InstrProfilingValue.c",
                            "InstrProfilingWriter.c"];

    for src in profile_sources {
        cfg.file(Path::new("../compiler-rt/lib/profile").join(src));
    }

    cfg.compile("libprofiler-rt.a");
}
<commit_msg>Compile WindowsMMap too on MSVC<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Compiles the profiler part of the `compiler-rt` library.
//!
//! See the build.rs for libcompiler_builtins crate for details.

extern crate gcc;

use std::env;
use std::path::Path;

fn main() {
    let target = env::var("TARGET").expect("TARGET was not set");
    let cfg = &mut gcc::Config::new();

    let mut profile_sources = vec!["GCDAProfiling.c",
                                   "InstrProfiling.c",
                                   "InstrProfilingBuffer.c",
                                   "InstrProfilingFile.c",
                                   "InstrProfilingMerge.c",
                                   "InstrProfilingMergeFile.c",
                                   "InstrProfilingPlatformDarwin.c",
                                   "InstrProfilingPlatformLinux.c",
                                   "InstrProfilingPlatformOther.c",
                                   "InstrProfilingRuntime.cc",
                                   "InstrProfilingUtil.c",
                                   "InstrProfilingValue.c",
                                   "InstrProfilingWriter.c"];

    if target.contains("msvc") {
        // Don't pull in extra libraries on MSVC
        cfg.flag("/Zl");
        profile_sources.push("WindowsMMap.c");
    } else {
        // Turn off various features of gcc and such, mostly copying
        // compiler-rt's build system already
        cfg.flag("-fno-builtin");
        cfg.flag("-fvisibility=hidden");
        cfg.flag("-fomit-frame-pointer");
        cfg.flag("-ffreestanding");
        cfg.define("VISIBILITY_HIDDEN", None);
    }

    for src in profile_sources {
        cfg.file(Path::new("../compiler-rt/lib/profile").join(src));
    }

    cfg.compile("libprofiler-rt.a");
}
<|endoftext|>


<commit_before>// Copyright 2014-2016 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Rank/Select data structure based on Gonzalez, Grabowski, Mäkinen, Navarro (2005).
//! This implementation uses only a single level of blocks, and performs well for large n.
//!
//! Example
//!
//! ```
//! extern crate bv;
//! # extern crate bio;
//! # fn main() {
//! use bio::data_structures::rank_select::RankSelect;
//! use bv::BitVec;
//! use bv::BitsMut;
//!
//! let mut bits: BitVec<u8> = BitVec::new_fill(false, 64);
//! bits.set_bit(5, true);
//! bits.set_bit(32, true);
//! let rs = RankSelect::new(bits, 1);
//! assert!(rs.rank(6).unwrap() == 1);
//! # }
//! ```

use std::cmp;

use bv::BitVec;
use bv::Bits;

/// A rank/select data structure.
#[derive(Serialize, Deserialize)]
pub struct RankSelect {
    n: usize,
    bits: BitVec<u8>,
    superblocks_1: Vec<u64>,
    superblocks_0: Vec<u64>,
    s: usize,
    k: usize,
}

impl RankSelect {
    /// Create a new instance.
    ///
    /// # Arguments
    ///
    /// * `bits` - A bit vector.
    /// * `k` - Determines the size (k * 32 bits) of the superblocks.
    ///   A small k means faster rank query times at the expense of using more
    ///   space and slower select query times.
    ///   The data structure needs O(n + n log n / (k * 32)) bits with n being the bits of the given bitvector.
    ///   The data structure is succinct if k is chosen as a sublinear function of n
    ///   (e.g. k = (log n)² / 32).
    pub fn new(bits: BitVec<u8>, k: usize) -> RankSelect {
        let n = bits.len() as usize;
        let s = k * 32;

        RankSelect {
            n,
            s,
            k,
            superblocks_1: superblocks(true, n, s, &bits),
            superblocks_0: superblocks(false, n, s, &bits),
            bits,
        }
    }

    /// Return the used k (see `RankSelect::new()`).
    pub fn k(&self) -> usize {
        self.k
    }

    /// Get internal representation of bit vector.
    pub fn bits(&self) -> &BitVec<u8> {
        &self.bits
    }

    /// Return i-th bit.
    pub fn get(&self, i: u64) -> bool {
        self.bits.get_bit(i)
    }

    /// Get the 1-rank of a given bit, i.e. the number of 1-bits in the bitvector up to i (inclusive).
    /// Complexity: O(k).
    ///
    /// # Arguments
    ///
    /// * `i` - Position of the bit to determine the rank for.
    pub fn rank_1(&self, i: u64) -> Option<u64> {
        if i >= self.n as u64 {
            None
        } else {
            let s = i / self.s as u64; // the superblock
            let b = i / 8; // the block
            let j = i % 8; // the bit in the block
                           // take the superblock rank
            let mut rank = self.superblocks_1[s as usize];
            // add the rank within the block
            rank += (self.bits.get_block(b as usize) & ((2 << j) - 1)).count_ones() as u64;
            // add the popcounts of blocks in between
            for block in s * 32 / 8..b {
                let b = self.bits.get_block(block as usize);
                rank += b.count_ones() as u64;
            }

            Some(rank)
        }
    }

    /// Get the 0-rank of a given bit, i.e. the number of 0-bits in the bitvector up to i (inclusive).
    /// Complexity: O(k).
    ///
    /// # Arguments
    ///
    /// * `i` - Position of the bit to determine the rank for.
    pub fn rank_0(&self, i: u64) -> Option<u64> {
        self.rank_1(i).map(|r| (i + 1) - r)
    }

    /// Alias for `RankSelect::rank_1`.
    pub fn rank(&self, i: u64) -> Option<u64> {
        self.rank_1(i)
    }

    /// Get the smallest bit with a given 1-rank.
    /// Complexity: O(log (n / k) + k).
    ///
    /// # Arguments
    ///
    /// * `j` - The rank to find the smallest bit for.
    pub fn select_1(&self, j: u64) -> Option<u64> {
        let mut superblock = match self.superblocks_1.binary_search(&j) {
            Ok(i) | Err(i) => i, // superblock with same rank exists
        };
        if superblock > 0 {
            superblock -= 1;
        }
        let mut rank = self.superblocks_1[superblock];

        let first_block = superblock * self.s / 8;
        for block in first_block..cmp::min(first_block + self.s / 8, self.bits.block_len()) {
            let b = self.bits.get_block(block);
            let p = b.count_ones() as u64;
            if rank + p >= j {
                let mut bit = 0b1;
                // do not look at unused bits of the last block
                let max_bit = cmp::min(8, self.bits.len() - block as u64 * 8);
                for i in 0..max_bit {
                    rank += (b & bit != 0) as u64;
                    if rank == j {
                        return Some((first_block + block) as u64 * 8 + i);
                    }
                    bit <<= 1;
                }
            }
            rank += p;
        }

        None
    }

    /// Get the smallest bit with a given 0-rank.
    /// Complexity: O(log (n / k) + k).
    ///
    /// # Arguments
    ///
    /// * `j` - The rank to find the smallest bit for.
    pub fn select_0(&self, j: u64) -> Option<u64> {
        let mut superblock = match self.superblocks_0.binary_search(&j) {
            Ok(i) | Err(i) => i, // superblock with same rank exists
        };
        if superblock > 0 {
            superblock -= 1;
        }
        let mut rank = self.superblocks_0[superblock];

        let first_block = superblock * self.s / 8;
        for block in first_block..cmp::min(first_block + self.s / 8, self.bits.block_len()) {
            let b = self.bits.get_block(block);
            let p = b.count_zeros() as u64;
            if rank + p >= j {
                let mut bit = 0b1;
                // do not look at unused bits of the last block
                let max_bit = cmp::min(8, self.bits.len() - block as u64 * 8);
                for i in 0..max_bit {
                    rank += (b & bit == 0) as u64;
                    if rank == j {
                        return Some((first_block + block) as u64 * 8 + i);
                    }
                    bit <<= 1;
                }
            }
            rank += p;
        }

        None
    }

    /// Alias for `RankSelect::select_1`.
    pub fn select(&self, j: u64) -> Option<u64> {
        self.select_1(j)
    }
}

/// Create `n` superblocks of size `s` from a given bitvector.
fn superblocks(t: bool, n: usize, s: usize, bits: &BitVec<u8>) -> Vec<u64> {
    let mut superblocks = Vec::with_capacity(n / s + 1);
    let mut rank: u64 = 0;
    let mut i = 0;
    let nblocks = (bits.len() as f64 / 8.0).ceil() as usize;
    for block in 0..nblocks {
        let b = bits.get_block(block);
        if i % s == 0 {
            superblocks.push(rank);
        }
        rank += if t {
            b.count_ones() as u64
        } else {
            b.count_zeros() as u64
        };
        i += 8;
    }

    superblocks
}

#[cfg(test)]
mod tests {
    use super::*;
    use bv::BitVec;
    use bv::BitsMut;

    #[test]
    fn test_rank_select() {
        let mut bits: BitVec<u8> = BitVec::new_fill(false, 64);
        bits.set_bit(5, true);
        bits.set_bit(32, true);
        let rs = RankSelect::new(bits, 1);
        assert_eq!(rs.rank_1(1).unwrap(), 0);
        assert_eq!(rs.rank_1(5).unwrap(), 1);
        assert_eq!(rs.rank_1(6).unwrap(), 1);
        assert_eq!(rs.rank_1(32).unwrap(), 2);
        assert_eq!(rs.rank_1(33).unwrap(), 2);
        assert_eq!(rs.rank_1(64), None);
        assert_eq!(rs.select_1(0).unwrap(), 0);
        assert_eq!(rs.select_1(1).unwrap(), 5);
        assert_eq!(rs.rank_0(1).unwrap(), 2);
        assert_eq!(rs.rank_0(4).unwrap(), 5);
        assert_eq!(rs.rank_0(5).unwrap(), 5);
        assert_eq!(rs.select_0(0), None);
        assert_eq!(rs.select_0(1).unwrap(), 0);
        assert_eq!(rs.get(5), true);
        assert_eq!(rs.get(1), false);
        assert_eq!(rs.get(32), true);
    }

    #[test]
    fn test_select() {
        let bits: BitVec<u8> = bit_vec![true, false];
        let rs = RankSelect::new(bits, 1);

        assert_eq!(rs.select_0(0), Some(0));
        assert_eq!(rs.select_1(0), None);

        assert_eq!(rs.select_0(1), Some(1));
        assert_eq!(rs.select_1(1), Some(0));

        assert_eq!(rs.select_0(2), None);
        assert_eq!(rs.select_1(2), None);
    }
}
<commit_msg>Remove redundant code in RankSelect. fix issue #172.<commit_after>// Copyright 2014-2016 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Rank/Select data structure based on Gonzalez, Grabowski, Mäkinen, Navarro (2005).
//! This implementation uses only a single level of blocks, and performs well for large n.
//!
//! Example
//!
//! ```
//! extern crate bv;
//! # extern crate bio;
//! # fn main() {
//! use bio::data_structures::rank_select::RankSelect;
//! use bv::BitVec;
//! use bv::BitsMut;
//!
//! let mut bits: BitVec<u8> = BitVec::new_fill(false, 64);
//! bits.set_bit(5, true);
//! bits.set_bit(32, true);
//! let rs = RankSelect::new(bits, 1);
//! assert!(rs.rank(6).unwrap() == 1);
//! # }
//! ```

use std::cmp;

use bv::BitVec;
use bv::Bits;

/// A rank/select data structure.
#[derive(Serialize, Deserialize)]
pub struct RankSelect {
    n: usize,
    bits: BitVec<u8>,
    superblocks_1: Vec<u64>,
    superblocks_0: Vec<u64>,
    s: usize,
    k: usize,
}

impl RankSelect {
    /// Create a new instance.
    ///
    /// # Arguments
    ///
    /// * `bits` - A bit vector.
    /// * `k` - Determines the size (k * 32 bits) of the superblocks.
    ///   A small k means faster rank query times at the expense of using more
    ///   space and slower select query times.
    ///   The data structure needs O(n + n log n / (k * 32)) bits with n being the bits of the given bitvector.
    ///   The data structure is succinct if k is chosen as a sublinear function of n
    ///   (e.g. k = (log n)² / 32).
    pub fn new(bits: BitVec<u8>, k: usize) -> RankSelect {
        let n = bits.len() as usize;
        let s = k * 32;

        RankSelect {
            n,
            s,
            k,
            superblocks_1: superblocks(true, n, s, &bits),
            superblocks_0: superblocks(false, n, s, &bits),
            bits,
        }
    }

    /// Return the used k (see `RankSelect::new()`).
    pub fn k(&self) -> usize {
        self.k
    }

    /// Get internal representation of bit vector.
    pub fn bits(&self) -> &BitVec<u8> {
        &self.bits
    }

    /// Return i-th bit.
    pub fn get(&self, i: u64) -> bool {
        self.bits.get_bit(i)
    }

    /// Get the 1-rank of a given bit, i.e. the number of 1-bits in the bitvector up to i (inclusive).
    /// Complexity: O(k).
    ///
    /// # Arguments
    ///
    /// * `i` - Position of the bit to determine the rank for.
    pub fn rank_1(&self, i: u64) -> Option<u64> {
        if i >= self.n as u64 {
            None
        } else {
            let s = i / self.s as u64; // the superblock
            let b = i / 8; // the block
            let j = i % 8; // the bit in the block
                           // take the superblock rank
            let mut rank = self.superblocks_1[s as usize];
            // add the rank within the block
            rank += (self.bits.get_block(b as usize) & ((2 << j) - 1)).count_ones() as u64;
            // add the popcounts of blocks in between
            for block in s * 32 / 8..b {
                let b = self.bits.get_block(block as usize);
                rank += b.count_ones() as u64;
            }

            Some(rank)
        }
    }

    /// Get the 0-rank of a given bit, i.e. the number of 0-bits in the bitvector up to i (inclusive).
    /// Complexity: O(k).
    ///
    /// # Arguments
    ///
    /// * `i` - Position of the bit to determine the rank for.
    pub fn rank_0(&self, i: u64) -> Option<u64> {
        self.rank_1(i).map(|r| (i + 1) - r)
    }

    /// Alias for `RankSelect::rank_1`.
    pub fn rank(&self, i: u64) -> Option<u64> {
        self.rank_1(i)
    }

    /// Get the smallest bit with a given 1-rank.
    /// Complexity: O(log (n / k) + k).
    ///
    /// # Arguments
    ///
    /// * `j` - The rank to find the smallest bit for.
    pub fn select_1(&self, j: u64) -> Option<u64> {
        self.select_x(j, &self.superblocks_1, |bit| bit != 0)
    }

    /// Get the smallest bit with a given 0-rank.
    /// Complexity: O(log (n / k) + k).
    ///
    /// # Arguments
    ///
    /// * `j` - The rank to find the smallest bit for.
    pub fn select_0(&self, j: u64) -> Option<u64> {
        self.select_x(j, &self.superblocks_0, |bit| bit == 0)
    }

    fn select_x<F: Fn(u8) -> bool>(&self, j: u64, superblocks: &[u64], is_match: F) -> Option<u64> {
        let mut superblock = match superblocks.binary_search(&j) {
            Ok(i) | Err(i) => i, // superblock with same rank exists
        };
        if superblock > 0 {
            superblock -= 1;
        }
        let mut rank = superblocks[superblock];

        let first_block = superblock * self.s / 8;
        for block in first_block..cmp::min(first_block + self.s / 8, self.bits.block_len()) {
            let b = self.bits.get_block(block);
            let p = b.count_zeros() as u64;
            if rank + p >= j {
                let mut bit = 0b1;
                // do not look at unused bits of the last block
                let max_bit = cmp::min(8, self.bits.len() - block as u64 * 8);
                for i in 0..max_bit {
                    rank += is_match(b & bit) as u64;
                    if rank == j {
                        return Some((first_block + block) as u64 * 8 + i);
                    }
                    bit <<= 1;
                }
            }
            rank += p;
        }

        None
    }

    /// Alias for `RankSelect::select_1`.
    pub fn select(&self, j: u64) -> Option<u64> {
        self.select_1(j)
    }
}

/// Create `n` superblocks of size `s` from a given bitvector.
fn superblocks(t: bool, n: usize, s: usize, bits: &BitVec<u8>) -> Vec<u64> {
    let mut superblocks = Vec::with_capacity(n / s + 1);
    let mut rank: u64 = 0;
    let mut i = 0;
    let nblocks = (bits.len() as f64 / 8.0).ceil() as usize;
    for block in 0..nblocks {
        let b = bits.get_block(block);
        if i % s == 0 {
            superblocks.push(rank);
        }
        rank += if t {
            b.count_ones() as u64
        } else {
            b.count_zeros() as u64
        };
        i += 8;
    }

    superblocks
}

#[cfg(test)]
mod tests {
    use super::*;
    use bv::BitVec;
    use bv::BitsMut;

    #[test]
    fn test_rank_select() {
        let mut bits: BitVec<u8> = BitVec::new_fill(false, 64);
        bits.set_bit(5, true);
        bits.set_bit(32, true);
        let rs = RankSelect::new(bits, 1);
        assert_eq!(rs.rank_1(1).unwrap(), 0);
        assert_eq!(rs.rank_1(5).unwrap(), 1);
        assert_eq!(rs.rank_1(6).unwrap(), 1);
        assert_eq!(rs.rank_1(32).unwrap(), 2);
        assert_eq!(rs.rank_1(33).unwrap(), 2);
        assert_eq!(rs.rank_1(64), None);
        assert_eq!(rs.select_1(0).unwrap(), 0);
        assert_eq!(rs.select_1(1).unwrap(), 5);
        assert_eq!(rs.rank_0(1).unwrap(), 2);
        assert_eq!(rs.rank_0(4).unwrap(), 5);
        assert_eq!(rs.rank_0(5).unwrap(), 5);
        assert_eq!(rs.select_0(0), None);
        assert_eq!(rs.select_0(1).unwrap(), 0);
        assert_eq!(rs.get(5), true);
        assert_eq!(rs.get(1), false);
        assert_eq!(rs.get(32), true);
    }

    #[test]
    fn test_select() {
        let bits: BitVec<u8> = bit_vec![true, false];
        let rs = RankSelect::new(bits, 1);

        assert_eq!(rs.select_0(0), Some(0));
        assert_eq!(rs.select_1(0), None);

        assert_eq!(rs.select_0(1), Some(1));
        assert_eq!(rs.select_1(1), Some(0));

        assert_eq!(rs.select_0(2), None);
        assert_eq!(rs.select_1(2), None);
    }
}
<|endoftext|>


<commit_before>use crate::spec::{LinkerFlavor, Target, TargetResult};

pub fn target() -> TargetResult {
    let mut base = super::freebsd_base::opts();
    base.cpu = "pentium4".to_string();
    base.max_atomic_width = Some(64);
    base.pre_link_args.get_mut(&LinkerFlavor::Gcc).unwrap().push("-m32".to_string());
    base.stack_probes = true;

    Ok(Target {
        llvm_target: "i686-unknown-freebsd".to_string(),
        target_endian: "little".to_string(),
        target_pointer_width: "32".to_string(),
        target_c_int_width: "32".to_string(),
        data_layout: "e-m:e-p:32:32-p270:32:32-p271:32:32-p272:64:64-\
            f64:32:64-f80:32-n8:16:32-S128"
            .to_string(),
        arch: "x86".to_string(),
        target_os: "freebsd".to_string(),
        target_env: String::new(),
        target_vendor: "unknown".to_string(),
        linker_flavor: LinkerFlavor::Gcc,
        options: base,
    })
}
<commit_msg>Add -Wl,-znotext to default linker flags to link with lld 9 on FreeBSD 13.0-CURRENT i386<commit_after>use crate::spec::{LinkerFlavor, Target, TargetResult};

pub fn target() -> TargetResult {
    let mut base = super::freebsd_base::opts();
    base.cpu = "pentium4".to_string();
    base.max_atomic_width = Some(64);
    let pre_link_args = base.pre_link_args.get_mut(&LinkerFlavor::Gcc).unwrap();
    pre_link_args.push("-m32".to_string());
    pre_link_args.push("-Wl,-znotext".to_string());
    base.stack_probes = true;

    Ok(Target {
        llvm_target: "i686-unknown-freebsd".to_string(),
        target_endian: "little".to_string(),
        target_pointer_width: "32".to_string(),
        target_c_int_width: "32".to_string(),
        data_layout: "e-m:e-p:32:32-p270:32:32-p271:32:32-p272:64:64-\
            f64:32:64-f80:32-n8:16:32-S128"
            .to_string(),
        arch: "x86".to_string(),
        target_os: "freebsd".to_string(),
        target_env: String::new(),
        target_vendor: "unknown".to_string(),
        linker_flavor: LinkerFlavor::Gcc,
        options: base,
    })
}
<|endoftext|>


<commit_before>use rustc_ast::token;
use rustc_data_structures::sync::{Lock, Lrc};
use rustc_errors::{emitter::Emitter, Applicability, Diagnostic, Handler};
use rustc_parse::lexer::StringReader as Lexer;
use rustc_session::parse::ParseSess;
use rustc_span::source_map::{FilePathMapping, SourceMap};
use rustc_span::{FileName, InnerSpan};

use crate::clean;
use crate::core::DocContext;
use crate::fold::DocFolder;
use crate::html::markdown::{self, RustCodeBlock};
use crate::passes::{span_of_attrs, Pass};

pub const CHECK_CODE_BLOCK_SYNTAX: Pass = Pass {
    name: "check-code-block-syntax",
    run: check_code_block_syntax,
    description: "validates syntax inside Rust code blocks",
};

pub fn check_code_block_syntax(krate: clean::Crate, cx: &DocContext<'_>) -> clean::Crate {
    SyntaxChecker { cx }.fold_crate(krate)
}

struct SyntaxChecker<'a, 'tcx> {
    cx: &'a DocContext<'tcx>,
}

impl<'a, 'tcx> SyntaxChecker<'a, 'tcx> {
    fn check_rust_syntax(&self, item: &clean::Item, dox: &str, code_block: RustCodeBlock) {
        let buffered_messages = Lrc::new(Lock::new(vec![]));

        let emitter = BufferEmitter { messages: Lrc::clone(&buffered_messages) };

        let sm = Lrc::new(SourceMap::new(FilePathMapping::empty()));
        let handler = Handler::with_emitter(false, None, Box::new(emitter));
        let sess = ParseSess::with_span_handler(handler, sm);
        let source_file = sess.source_map().new_source_file(
            FileName::Custom(String::from("doctest")),
            dox[code_block.code].to_owned(),
        );

        let validation_status = rustc_driver::catch_fatal_errors(|| {
            let mut has_syntax_errors = false;
            let mut only_whitespace = true;
            // even if there is a syntax error, we need to run the lexer over the whole file
            let mut lexer = Lexer::new(&sess, source_file, None);
            loop {
                match lexer.next_token().kind {
                    token::Eof => break,
                    token::Whitespace => (),
                    token::Unknown(..) => has_syntax_errors = true,
                    _ => only_whitespace = false,
                }
            }

            if has_syntax_errors {
                Some(CodeBlockInvalid::SyntaxError)
            } else if only_whitespace {
                Some(CodeBlockInvalid::Empty)
            } else {
                None
            }
        })
        .unwrap_or(Some(CodeBlockInvalid::SyntaxError));

        if let Some(code_block_invalid) = validation_status {
            let mut diag = if let Some(sp) =
                super::source_span_for_markdown_range(self.cx, &dox, &code_block.range, &item.attrs)
            {
                let warning_message = match code_block_invalid {
                    CodeBlockInvalid::SyntaxError => "could not parse code block as Rust code",
                    CodeBlockInvalid::Empty => "Rust code block is empty",
                };

                let mut diag = self.cx.sess().struct_span_warn(sp, warning_message);

                if code_block.syntax.is_none() && code_block.is_fenced {
                    let sp = sp.from_inner(InnerSpan::new(0, 3));
                    diag.span_suggestion(
                        sp,
                        "mark blocks that do not contain Rust code as text",
                        String::from("```text"),
                        Applicability::MachineApplicable,
                    );
                }

                diag
            } else {
                // We couldn't calculate the span of the markdown block that had the error, so our
                // diagnostics are going to be a bit lacking.
                let mut diag = self.cx.sess().struct_span_warn(
                    super::span_of_attrs(&item.attrs).unwrap_or(item.source.span()),
                    "doc comment contains an invalid Rust code block",
                );

                if code_block.syntax.is_none() && code_block.is_fenced {
                    diag.help("mark blocks that do not contain Rust code as text: ```text");
                }

                diag
            };

            // FIXME(#67563): Provide more context for these errors by displaying the spans inline.
            for message in buffered_messages.borrow().iter() {
                diag.note(&message);
            }

            diag.emit();
        }
    }
}

impl<'a, 'tcx> DocFolder for SyntaxChecker<'a, 'tcx> {
    fn fold_item(&mut self, item: clean::Item) -> Option<clean::Item> {
        if let Some(dox) = &item.attrs.collapsed_doc_value() {
            let sp = span_of_attrs(&item.attrs).unwrap_or(item.source.span());
            let extra = crate::html::markdown::ExtraInfo::new_did(&self.cx.tcx, item.def_id, sp);
            for code_block in markdown::rust_code_blocks(&dox, &extra) {
                self.check_rust_syntax(&item, &dox, code_block);
            }
        }

        self.fold_item_recur(item)
    }
}

struct BufferEmitter {
    messages: Lrc<Lock<Vec<String>>>,
}

impl Emitter for BufferEmitter {
    fn emit_diagnostic(&mut self, diag: &Diagnostic) {
        self.messages.borrow_mut().push(format!("error from rustc: {}", diag.message[0].0));
    }

    fn source_map(&self) -> Option<&Lrc<SourceMap>> {
        None
    }
}

enum CodeBlockInvalid {
    SyntaxError,
    Empty,
}
<commit_msg>Avoid StringReader when checking code blocks for syntax errors<commit_after>use rustc_data_structures::sync::{Lock, Lrc};
use rustc_errors::{emitter::Emitter, Applicability, Diagnostic, Handler};
use rustc_parse::parse_stream_from_source_str;
use rustc_session::parse::ParseSess;
use rustc_span::source_map::{FilePathMapping, SourceMap};
use rustc_span::{FileName, InnerSpan};

use crate::clean;
use crate::core::DocContext;
use crate::fold::DocFolder;
use crate::html::markdown::{self, RustCodeBlock};
use crate::passes::{span_of_attrs, Pass};

pub const CHECK_CODE_BLOCK_SYNTAX: Pass = Pass {
    name: "check-code-block-syntax",
    run: check_code_block_syntax,
    description: "validates syntax inside Rust code blocks",
};

pub fn check_code_block_syntax(krate: clean::Crate, cx: &DocContext<'_>) -> clean::Crate {
    SyntaxChecker { cx }.fold_crate(krate)
}

struct SyntaxChecker<'a, 'tcx> {
    cx: &'a DocContext<'tcx>,
}

impl<'a, 'tcx> SyntaxChecker<'a, 'tcx> {
    fn check_rust_syntax(&self, item: &clean::Item, dox: &str, code_block: RustCodeBlock) {
        let buffer = Lrc::new(Lock::new(Buffer::default()));
        let emitter = BufferEmitter { buffer: Lrc::clone(&buffer) };

        let sm = Lrc::new(SourceMap::new(FilePathMapping::empty()));
        let handler = Handler::with_emitter(false, None, Box::new(emitter));
        let source = dox[code_block.code].to_owned();
        let sess = ParseSess::with_span_handler(handler, sm);

        let is_empty = rustc_driver::catch_fatal_errors(|| {
            parse_stream_from_source_str(
                FileName::Custom(String::from("doctest")),
                source,
                &sess,
                None,
            )
            .is_empty()
        })
        .unwrap_or(false);
        let buffer = buffer.borrow();

        if buffer.has_errors || is_empty {
            let mut diag = if let Some(sp) =
                super::source_span_for_markdown_range(self.cx, &dox, &code_block.range, &item.attrs)
            {
                let warning_message = if buffer.has_errors {
                    "could not parse code block as Rust code"
                } else {
                    "Rust code block is empty"
                };

                let mut diag = self.cx.sess().struct_span_warn(sp, warning_message);

                if code_block.syntax.is_none() && code_block.is_fenced {
                    let sp = sp.from_inner(InnerSpan::new(0, 3));
                    diag.span_suggestion(
                        sp,
                        "mark blocks that do not contain Rust code as text",
                        String::from("```text"),
                        Applicability::MachineApplicable,
                    );
                }

                diag
            } else {
                // We couldn't calculate the span of the markdown block that had the error, so our
                // diagnostics are going to be a bit lacking.
                let mut diag = self.cx.sess().struct_span_warn(
                    super::span_of_attrs(&item.attrs).unwrap_or(item.source.span()),
                    "doc comment contains an invalid Rust code block",
                );

                if code_block.syntax.is_none() && code_block.is_fenced {
                    diag.help("mark blocks that do not contain Rust code as text: ```text");
                }

                diag
            };

            // FIXME(#67563): Provide more context for these errors by displaying the spans inline.
            for message in buffer.messages.iter() {
                diag.note(&message);
            }

            diag.emit();
        }
    }
}

impl<'a, 'tcx> DocFolder for SyntaxChecker<'a, 'tcx> {
    fn fold_item(&mut self, item: clean::Item) -> Option<clean::Item> {
        if let Some(dox) = &item.attrs.collapsed_doc_value() {
            let sp = span_of_attrs(&item.attrs).unwrap_or(item.source.span());
            let extra = crate::html::markdown::ExtraInfo::new_did(&self.cx.tcx, item.def_id, sp);
            for code_block in markdown::rust_code_blocks(&dox, &extra) {
                self.check_rust_syntax(&item, &dox, code_block);
            }
        }

        self.fold_item_recur(item)
    }
}

#[derive(Default)]
struct Buffer {
    messages: Vec<String>,
    has_errors: bool,
}

struct BufferEmitter {
    buffer: Lrc<Lock<Buffer>>,
}

impl Emitter for BufferEmitter {
    fn emit_diagnostic(&mut self, diag: &Diagnostic) {
        let mut buffer = self.buffer.borrow_mut();
        buffer.messages.push(format!("error from rustc: {}", diag.message[0].0));
        if diag.is_error() {
            buffer.has_errors = true;
        }
    }

    fn source_map(&self) -> Option<&Lrc<SourceMap>> {
        None
    }
}
<|endoftext|>


<commit_before><commit_msg>Make DecayOptions useful<commit_after><|endoftext|>


<commit_before><commit_msg>Stop allowing patterns to be created implicitly<commit_after><|endoftext|>


<commit_before><commit_msg>add filter policy interface<commit_after><|endoftext|>


<commit_before><commit_msg>Mac args correctly filtered.<commit_after><|endoftext|>


<commit_before><commit_msg>Add PB0 & PB1 in ADC.<commit_after><|endoftext|>


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

//! Provides wrappers for throttling readers and writers

use std::time::{Instant, Duration};
use std::io;

use futures::*;
use tokio_io::*;
use bytes::{Buf, BytesMut, BufMut};

/// A Rate Limited Reader
#[derive(Debug)]
pub struct ThrottledReader<R: AsyncRead> {
	reader: R,
	/// Max Bytes per second
	max: u32,
	/// Stores a count of last request and last request time
	allowed: isize,
	last_check: Instant,
}

impl<R: AsyncRead> ThrottledReader<R> {
	/// Adds throttling to a reader.
	/// The resulting reader will read at most `max` amount of bytes per second
	pub fn new(reader: R, max: u32) -> Self {
		ThrottledReader {
			reader: reader,
			max: max,
			allowed: max as isize,
			last_check: Instant::now(),
		}
	}

	/// Get a shared reference to the inner sink.
	pub fn get_ref(&self) -> &R {
		&self.reader
	}

	/// Get a mutable reference to the inner sink.
	pub fn get_mut(&mut self) -> &mut R {
		&mut self.reader
	}

	/// Consumes this combinator, returning the underlying sink.
	///
	/// Note that this may discard intermediate state of this combinator, so
	/// care should be taken to avoid losing resources when this is called.
	pub fn into_inner(self) -> R {
		self.reader
	}
}

impl<R: AsyncRead> io::Read for ThrottledReader<R> {
	fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
		self.reader.read(buf)
	}
}

impl<R: AsyncRead> AsyncRead for ThrottledReader<R> {
	unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
		self.reader.prepare_uninitialized_buffer(buf)
	}

	fn read_buf<B: BufMut>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
		// Check passed Time
		let time_passed = self.last_check.elapsed();
		self.last_check = Instant::now();
		self.allowed += time_passed.as_secs() as isize * self.max as isize;

		// Throttle
		if self.allowed > self.max as isize {
			self.allowed = self.max as isize;
		}

		// Check if Allowed
		if self.allowed < 1 {
			return Ok(Async::NotReady);
		}

		// Since we can't limit the scope that is read,
		// we use a signed `allowed` counter in case n > allowed
		let res = self.reader.read_buf(buf);

		// Decrement Allowed amount written
		if let Ok(Async::Ready(n)) = res {
			self.allowed = self.allowed.saturating_sub(n as isize);
		}
		res
	}
}

/// A Rate Limited Writer
#[derive(Debug)]
pub struct ThrottledWriter<W: AsyncWrite> {
	writer: W,
	/// Max Bytes per second
	max: u32,
	/// Stores a count of last request and last request time
	allowed: usize,
	last_check: Instant,
}

impl<W: AsyncWrite> ThrottledWriter<W> {
	/// Adds throttling to a writer.
	/// The resulting writer will write at most `max` amount of bytes per second
	pub fn new(writer: W, max: u32) -> Self {
		ThrottledWriter {
			writer: writer,
			max: max,
			allowed: max as usize,
			last_check: Instant::now(),
		}
	}

	/// Get a shared reference to the inner sink.
	pub fn get_ref(&self) -> &W {
		&self.writer
	}

	/// Get a mutable reference to the inner sink.
	pub fn get_mut(&mut self) -> &mut W {
		&mut self.writer
	}

	/// Consumes this combinator, returning the underlying sink.
	///
	/// Note that this may discard intermediate state of this combinator, so
	/// care should be taken to avoid losing resources when this is called.
	pub fn into_inner(self) -> W {
		self.writer
	}
}

impl<W: AsyncWrite> io::Write for ThrottledWriter<W> {
	fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
		self.writer.write(buf)
	}

	fn flush(&mut self) -> io::Result<()> {
		self.writer.flush()
	}
}

impl<T: AsyncWrite> AsyncWrite for ThrottledWriter<T> {
	fn shutdown(&mut self) -> Poll<(), io::Error> {
		self.writer.shutdown()
	}

	fn write_buf<B: Buf>(&mut self, buf: &mut B) -> Poll<usize, io::Error>
		where Self: Sized
	{
		// Check passed Time
		let time_passed = self.last_check.elapsed();
		self.last_check = Instant::now();
		self.allowed += time_passed.as_secs() as usize * self.max as usize;

		// Throttle
		if self.allowed > self.max as usize {
			self.allowed = self.max as usize;
		}

		// Check if Allowed
		if self.allowed < 1 {
			return Ok(Async::NotReady);
		}

		// Write max allowed
		let ref mut lbuf = buf.by_ref().take(self.allowed);
		let res = self.writer.write_buf(lbuf);

		// Decrement Allowed amount written
		if let Ok(Async::Ready(n)) = res {
			self.allowed -= n;
		}
		res
	}
}
<commit_msg>Use io::Read/Write interface instead of AsyncRead/AsyncWrite + Finish Tests (#74)<commit_after>// Copyright 2016 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Provides wrappers for throttling readers and writers

use std::time::{Instant, Duration};
use std::io;

use futures::*;
use tokio_io::*;
use bytes::{Buf, BytesMut, BufMut};

/// A Rate Limited Reader
#[derive(Debug)]
pub struct ThrottledReader<R: AsyncRead> {
	reader: R,
	/// Max Bytes per second
	max: u32,
	/// Stores a count of last request and last request time
	allowed: usize,
	last_check: Instant,
}

impl<R: AsyncRead> ThrottledReader<R> {
	/// Adds throttling to a reader.
	/// The resulting reader will read at most `max` amount of bytes per second
	pub fn new(reader: R, max: u32) -> Self {
		ThrottledReader {
			reader: reader,
			max: max,
			allowed: max as usize,
			last_check: Instant::now(),
		}
	}

	/// Get a shared reference to the inner sink.
	pub fn get_ref(&self) -> &R {
		&self.reader
	}

	/// Get a mutable reference to the inner sink.
	pub fn get_mut(&mut self) -> &mut R {
		&mut self.reader
	}

	/// Consumes this combinator, returning the underlying sink.
	///
	/// Note that this may discard intermediate state of this combinator, so
	/// care should be taken to avoid losing resources when this is called.
	pub fn into_inner(self) -> R {
		self.reader
	}
}

impl<R: AsyncRead> io::Read for ThrottledReader<R> {
	fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
		// Check passed Time
		let time_passed = self.last_check.elapsed();
		self.last_check = Instant::now();
		self.allowed += time_passed.as_secs() as usize * self.max as usize;

		// Throttle
		if self.allowed > self.max as usize {
			self.allowed = self.max as usize;
		}

		// Check if Allowed
		if self.allowed < 1 {
			return Err(io::Error::new(io::ErrorKind::WouldBlock, "Reached Allowed Read Limit"))
		}

		// Read Max Allowed
		let buf = if buf.len() > self.allowed { &mut buf[0..self.allowed]} else { buf };
		let res = self.reader.read(buf);

		// Decrement Allowed amount written
		if let Ok(n) = res {
			self.allowed -= n;
		}
		res
	}
}

impl<R: AsyncRead> AsyncRead for ThrottledReader<R> { }

/// A Rate Limited Writer
#[derive(Debug)]
pub struct ThrottledWriter<W: AsyncWrite> {
	writer: W,
	/// Max Bytes per second
	max: u32,
	/// Stores a count of last request and last request time
	allowed: usize,
	last_check: Instant,
}

impl<W: AsyncWrite> ThrottledWriter<W> {
	/// Adds throttling to a writer.
	/// The resulting writer will write at most `max` amount of bytes per second
	pub fn new(writer: W, max: u32) -> Self {
		ThrottledWriter {
			writer: writer,
			max: max,
			allowed: max as usize,
			last_check: Instant::now(),
		}
	}

	/// Get a shared reference to the inner sink.
	pub fn get_ref(&self) -> &W {
		&self.writer
	}

	/// Get a mutable reference to the inner sink.
	pub fn get_mut(&mut self) -> &mut W {
		&mut self.writer
	}

	/// Consumes this combinator, returning the underlying sink.
	///
	/// Note that this may discard intermediate state of this combinator, so
	/// care should be taken to avoid losing resources when this is called.
	pub fn into_inner(self) -> W {
		self.writer
	}
}

impl<W: AsyncWrite> io::Write for ThrottledWriter<W> {
	fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
		// Check passed Time
		let time_passed = self.last_check.elapsed();
		self.last_check = Instant::now();
		self.allowed += time_passed.as_secs() as usize * self.max as usize;

		// Throttle
		if self.allowed > self.max as usize {
			self.allowed = self.max as usize;
		}

		// Check if Allowed
		if self.allowed < 1 {
			return Err(io::Error::new(io::ErrorKind::WouldBlock, "Reached Allowed Write Limit"))
		}

		// Write max allowed
		let buf = if buf.len() > self.allowed { &buf[0..self.allowed]} else { buf };
		let res = self.writer.write(buf);

		// Decrement Allowed amount written
		if let Ok(n) = res {
			self.allowed -= n;
		}
		res
	}

	fn flush(&mut self) -> io::Result<()> {
		self.writer.flush()
	}
}

impl<T: AsyncWrite> AsyncWrite for ThrottledWriter<T> {
	fn shutdown(&mut self) -> Poll<(), io::Error> {
		self.writer.shutdown()
	}
}

#[cfg(test)]
mod test {
	use super::*;
	use std::io::Cursor;

	#[test]
	fn should_throttle_write() {
		let buf = vec![0; 64];
		let mut t_buf = ThrottledWriter::new(Cursor::new(buf), 8);

		for _ in 0..16 {
			let _ = t_buf.write_buf(&mut Cursor::new(vec![1; 8]));
		}

		let cursor = t_buf.into_inner();
		assert_eq!(cursor.position(), 8);
	}

	#[test]
	fn should_throttle_read() {
		let buf = vec![1; 64];
		let mut t_buf = ThrottledReader::new(Cursor::new(buf), 8);

		let mut dst = Cursor::new(vec![0; 64]);

		for _ in 0..16 {
			let _ = t_buf.read_buf(&mut dst);
		}

		assert_eq!(dst.position(), 8);
	}
}<|endoftext|>