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

#![allow(clippy::doc_markdown)]

#[macro_use]
extern crate log;

use std::{
    cell::RefCell,
    env,
    fs::{File, OpenOptions},
    io::{Read, Write},
    os::unix::io::{AsRawFd, RawFd},
    process::exit,
    rc::Rc,
};

use chrono::Duration;
use clap::{App, Arg, ArgMatches};
use env_logger::Builder;
use libc::pid_t;
use log::LevelFilter;
use nix::{
    fcntl::{flock, FlockArg},
    sys::{
        signal::{self, SigSet},
        signalfd::{SfdFlags, SignalFd},
    },
    unistd::getpid,
};
use uuid::Uuid;

#[cfg(feature = "dbus_enabled")]
use libstratis::{
    dbus_api::{DbusConnectionData, EventHandler},
    engine::get_engine_listener_list_mut,
};

use libstratis::{
    engine::{Engine, Pool, SimEngine, StratEngine},
    stratis::{buff_log, StratisError, StratisResult, VERSION},
};

const STRATISD_PID_PATH: &str = "/run/stratisd.pid";

/// Number of minutes to buffer log entries.
const DEFAULT_LOG_HOLD_MINUTES: i64 = 30;

/// If writing a program error to stderr fails, panic.
fn print_err(err: &StratisError) {
    eprintln!("{}", err);
}

/// Configure the env_logger as necessary in order to allow the buffered
/// logger to work correctly. Return a Handle to the underlying env_logger.
pub fn from_env_logger(
    mut builder: env_logger::Builder,
    pass_through: bool,
    hold_time: Option<Duration>,
) -> buff_log::Handle<env_logger::Logger> {
    // Do not have the env_logger set the timestamp. Because the entries are
    // buffered, the timestamp set by the env_logger will correspond to the
    // time at which the entry was dumped, not the time of its origination.
    builder.default_format_timestamp(false);
    buff_log::Logger::new(builder.build(), pass_through, hold_time).init()
}

/// Configure and initialize the logger.
/// If debug is true, log at debug level. Otherwise read log configuration
/// parameters from the environment if RUST_LOG is set. Otherwise, just
/// accept the default configuration.
fn initialize_log(debug: bool) -> buff_log::Handle<env_logger::Logger> {
    let mut builder = Builder::new();
    if debug {
        builder.filter(Some("stratisd"), LevelFilter::Debug);
        builder.filter(Some("libstratis"), LevelFilter::Debug);
        from_env_logger(builder, true, None)
    } else {
        builder.filter_level(LevelFilter::Trace);
        if let Ok(s) = env::var("RUST_LOG") {
            builder.parse(&s);
        }
        from_env_logger(
            builder,
            false,
            Some(Duration::minutes(DEFAULT_LOG_HOLD_MINUTES)),
        )
    }
}

/// To ensure only one instance of stratisd runs at a time, acquire an
/// exclusive lock. Return an error if lock attempt fails.
fn trylock_pid_file() -> StratisResult<File> {
    let mut f = OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .open(STRATISD_PID_PATH)
        .map_err(|err| {
            StratisError::Error(format!(
                "Failed to create or open the stratisd PID file at {}: {}",
                STRATISD_PID_PATH, err
            ))
        })?;
    match flock(f.as_raw_fd(), FlockArg::LockExclusiveNonblock) {
        Ok(_) => {
            f.write_all(format!("{}\n", getpid()).as_bytes())?;
            Ok(f)
        }
        Err(_) => {
            let mut buf = String::new();
            f.read_to_string(&mut buf)?;

            let pid_str = buf
                .split_whitespace()
                .next()
                .and_then(|s| s.parse::<pid_t>().ok())
                .map(|pid| format!("{}", pid))
                .unwrap_or_else(|| "<unknown>".into());
            Err(StratisError::Error(format!(
                "Daemon already running with pid: {}",
                pid_str
            )))
        }
    }
}

// Conditionally compiled support for a D-Bus interface.
struct MaybeDbusSupport {
    #[cfg(feature = "dbus_enabled")]
    handle: Option<libstratis::dbus_api::DbusConnectionData>,
}

// If D-Bus compiled out, do very little.
#[cfg(not(feature = "dbus_enabled"))]
impl MaybeDbusSupport {
    fn new() -> MaybeDbusSupport {
        MaybeDbusSupport {}
    }

    fn process(
        &mut self,
        _engine: &Rc<RefCell<dyn Engine>>,
        _fds: &mut Vec<libc::pollfd>,
        _dbus_client_index_start: usize,
    ) {
    }

    fn register_pool(&mut self, _pool_uuid: Uuid, _pool: &mut dyn Pool) {}

    fn poll_timeout(&self) -> i32 {
        // Non-DBus timeout is infinite
        -1
    }
}

#[cfg(feature = "dbus_enabled")]
impl MaybeDbusSupport {
    fn new() -> MaybeDbusSupport {
        MaybeDbusSupport { handle: None }
    }

    /// Connect to D-Bus and register pools, if not already connected.
    /// Return the connection, if made or already existing, otherwise, None.
    fn setup_connection(
        &mut self,
        engine: &Rc<RefCell<dyn Engine>>,
    ) -> Option<&mut DbusConnectionData> {
        if self.handle.is_none() {
            match libstratis::dbus_api::DbusConnectionData::connect(Rc::clone(engine)) {
                Err(err) => {
                    warn!("D-Bus API is not available: {}", err);
                }
                Ok(mut handle) => {
                    info!("D-Bus API is available");
                    let event_handler = Box::new(EventHandler::new(Rc::clone(&handle.connection)));
                    get_engine_listener_list_mut().register_listener(event_handler);
                    // Register all the pools with dbus
                    for (_, pool_uuid, pool) in engine.borrow_mut().pools_mut() {
                        handle.register_pool(pool_uuid, pool)
                    }
                    self.handle = Some(handle);
                }
            }
        };
        self.handle.as_mut()
    }

    /// Handle any client dbus requests.
    fn process(
        &mut self,
        engine: &Rc<RefCell<dyn Engine>>,
        fds: &mut Vec<libc::pollfd>,
        dbus_client_index_start: usize,
    ) {
        if let Some(handle) = self.setup_connection(engine) {
            handle.handle(&fds[dbus_client_index_start..]);

            // Refresh list of dbus fds to poll for. This can change as
            // D-Bus clients come and go.
            fds.truncate(dbus_client_index_start);
            fds.extend(
                handle
                    .connection
                    .borrow()
                    .watch_fds()
                    .iter()
                    .map(|w| w.to_pollfd()),
            );
        }
    }

    fn register_pool(&mut self, pool_uuid: Uuid, pool: &mut dyn Pool) {
        if let Some(h) = self.handle.as_mut() {
            h.register_pool(pool_uuid, pool)
        }
    }

    fn poll_timeout(&self) -> i32 {
        // If there is no D-Bus connection set timeout to 1 sec (1000 ms), so
        // that stratisd can periodically attempt to set up a connection.
        // If the connection is up, set the timeout to infinite; there is no
        // need to poll as events will be received.
        self.handle.as_ref().map_or(1000, |_| -1)
    }
}

// A facility for listening for and handling udev events that stratisd
// considers interesting.
struct UdevMonitor<'a> {
    socket: libudev::MonitorSocket<'a>,
}

impl<'a> UdevMonitor<'a> {
    fn create(context: &'a libudev::Context) -> StratisResult<UdevMonitor<'a>> {
        let mut monitor = libudev::Monitor::new(context)?;
        monitor.match_subsystem("block")?;

        Ok(UdevMonitor {
            socket: monitor.listen()?,
        })
    }

    fn as_raw_fd(&mut self) -> RawFd {
        self.socket.as_raw_fd()
    }

    /// Handle udev events.
    /// Check if a pool can be constructed and update engine and D-Bus layer
    /// data structures if so.
    fn handle_events(&mut self, engine: &mut dyn Engine, dbus_support: &mut MaybeDbusSupport) {
        while let Some(event) = self.socket.receive_event() {
            if let Some((pool_uuid, pool)) = engine.handle_event(&event) {
                dbus_support.register_pool(pool_uuid, pool);
            }
        }
    }
}

// Process any pending signals, return true if SIGINT received.
// Return an error if there was an error reading the signal.
fn process_signal(
    sfd: &mut SignalFd,
    buff_log: &buff_log::Handle<env_logger::Logger>,
) -> StratisResult<bool> {
    match sfd.read_signal() {
        // This is an unsafe conversion, but in this context that is
        // mostly harmless. A negative converted value, which is
        // virtually impossible, will not match any of the masked
        // values, and stratisd will panic and exit.
        Ok(Some(sig)) => match sig.ssi_signo as i32 {
            nix::libc::SIGUSR1 => {
                info!(
                    "SIGUSR1 received, dumping {} buffered log entries",
                    buff_log.buffered_count()
                );
                buff_log.dump();
                Ok(false)
            }
            nix::libc::SIGINT => {
                info!("SIGINT received, exiting");
                Ok(true)
            }
            signo => {
                panic!("Caught an impossible signal {:?}", signo);
            }
        },
        // No signals waiting (SFD_NONBLOCK flag is set)
        Ok(None) => Ok(false),

        Err(err) => Err(err.into()),
    }
}

/// Handle blocking the event loop
fn process_poll(poll_timeout: i32, fds: &mut Vec<libc::pollfd>) -> StratisResult<()> {
    let r = unsafe { libc::poll(fds.as_mut_ptr(), fds.len() as libc::c_ulong, poll_timeout) };

    // TODO: refine this behavior.
    // Different behaviors may be indicated, depending on the value of
    // errno when return value is -1.
    if r < 0 {
        return Err(StratisError::Error(format!(
            "poll command failed: number of fds: {}, timeout: {}",
            fds.len(),
            poll_timeout
        )));
    }
    Ok(())
}

/// Set up all sorts of signal and event handling mechanisms.
/// Initialize the engine and keep it running until a signal is received
/// or a fatal error is encountered. Dump log entries on specified signal
/// via buff_log.
fn run(matches: &ArgMatches, buff_log: &buff_log::Handle<env_logger::Logger>) -> StratisResult<()> {
    // Ensure that the debug log is output when we leave this function.
    let _guard = buff_log.to_guard();

    let mut dbus_support = MaybeDbusSupport::new();

    // Setup a udev listener before initializing the engine. A device may
    // appear after the engine has processed the udev db, but before it has
    // completed initialization. Unless the udev event has been recorded, the
    // engine will miss the device.
    // This is especially important since stratisd must run during early boot.
    let context = libudev::Context::new()?;
    let mut udev_monitor = UdevMonitor::create(&context)?;

    let engine: Rc<RefCell<dyn Engine>> = {
        info!("stratis daemon version {} started", VERSION);
        if matches.is_present("sim") {
            info!("Using SimEngine");
            Rc::new(RefCell::new(SimEngine::default()))
        } else {
            info!("Using StratEngine");
            Rc::new(RefCell::new(StratEngine::initialize()?))
        }
    };

    /*
    The file descriptor array indexes are:

    0   == Always udev fd index
    1   == SIGNAL FD index
    2   == engine index if eventable
    2/3 == Start of dbus client file descriptor(s)
            * 2 if engine is not eventable
            * else 3
    */
    const FD_INDEX_UDEV: usize = 0;
    const FD_INDEX_SIGNALFD: usize = 1;
    const FD_INDEX_ENGINE: usize = 2;

    /*
    fds is a Vec of libc::pollfd structs. Ideally, it would be possible
    to use the higher level nix crate to handle polling. If this were possible,
    then the Vec would be one of nix::poll::PollFds and this would be more
    rustic. Unfortunately, the rust D-Bus library requires an explicit file
    descriptor to be passed as an argument to Connection::watch_handle(),
    and the explicit file descriptor can not be extracted from the PollFd
    struct. So, at this time, sticking with libc is less complex than
    converting to using nix, because if using nix, the file descriptor would
    have to be maintained in the Vec as well as the PollFd struct.
    */
    let mut fds = Vec::new();

    fds.push(libc::pollfd {
        fd: udev_monitor.as_raw_fd(),
        revents: 0,
        events: libc::POLLIN,
    });

    // Signals can be queued up on this file descriptor
    let mut sfd = {
        let mut mask = SigSet::empty();
        mask.add(signal::SIGINT);
        mask.add(signal::SIGUSR1);
        mask.thread_block()?;
        SignalFd::with_flags(&mask, SfdFlags::SFD_NONBLOCK)?
    };

    fds.push(libc::pollfd {
        fd: sfd.as_raw_fd(),
        revents: 0,
        events: libc::POLLIN,
    });

    let eventable = engine.borrow().get_eventable();

    if let Some(evt) = eventable {
        fds.push(libc::pollfd {
            fd: evt.get_pollable_fd(),
            revents: 0,
            events: libc::POLLIN,
        });
    };

    let dbus_client_index_start = if eventable.is_some() {
        FD_INDEX_ENGINE + 1
    } else {
        FD_INDEX_ENGINE
    };

    loop {
        if fds[FD_INDEX_UDEV].revents != 0 {
            udev_monitor.handle_events(&mut *engine.borrow_mut(), &mut dbus_support)
        }

        if fds[FD_INDEX_SIGNALFD].revents != 0 {
            match process_signal(&mut sfd, buff_log) {
                Ok(should_exit) => {
                    if should_exit {
                        return Ok(());
                    }
                }
                Err(e) => return Err(e),
            }
        }

        if let Some(evt) = eventable {
            if fds[FD_INDEX_ENGINE].revents != 0 {
                evt.clear_event()?;
                engine.borrow_mut().evented()?;
            }
        }

        dbus_support.process(&engine, &mut fds, dbus_client_index_start);

        process_poll(dbus_support.poll_timeout(), &mut fds)?;
    }
}

fn main() {
    let matches = App::new("stratis")
        .version(VERSION)
        .about("Stratis storage management")
        .arg(
            Arg::with_name("debug")
                .long("debug")
                .help("Print additional output for debugging"),
        )
        .arg(
            Arg::with_name("sim")
                .long("sim")
                .help("Use simulator engine"),
        )
        .get_matches();

    // Using a let-expression here so that the scope of the lock file
    // is the rest of the block.
    let lock_file = trylock_pid_file();

    let result = {
        match lock_file {
            Err(err) => Err(err),
            Ok(_) => {
                let log_handle = initialize_log(matches.is_present("debug"));
                run(&matches, &log_handle)
            }
        }
    };

    if let Err(err) = result {
        print_err(&err);
        exit(1);
    } else {
        exit(0);
    }
}
<commit_msg>In error message, just print out what was read from the PID file<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/.

#![allow(clippy::doc_markdown)]

#[macro_use]
extern crate log;

use std::{
    cell::RefCell,
    env,
    fs::{File, OpenOptions},
    io::{Read, Write},
    os::unix::io::{AsRawFd, RawFd},
    process::exit,
    rc::Rc,
};

use chrono::Duration;
use clap::{App, Arg, ArgMatches};
use env_logger::Builder;
use log::LevelFilter;
use nix::{
    fcntl::{flock, FlockArg},
    sys::{
        signal::{self, SigSet},
        signalfd::{SfdFlags, SignalFd},
    },
    unistd::getpid,
};
use uuid::Uuid;

#[cfg(feature = "dbus_enabled")]
use libstratis::{
    dbus_api::{DbusConnectionData, EventHandler},
    engine::get_engine_listener_list_mut,
};

use libstratis::{
    engine::{Engine, Pool, SimEngine, StratEngine},
    stratis::{buff_log, StratisError, StratisResult, VERSION},
};

const STRATISD_PID_PATH: &str = "/run/stratisd.pid";

/// Number of minutes to buffer log entries.
const DEFAULT_LOG_HOLD_MINUTES: i64 = 30;

/// If writing a program error to stderr fails, panic.
fn print_err(err: &StratisError) {
    eprintln!("{}", err);
}

/// Configure the env_logger as necessary in order to allow the buffered
/// logger to work correctly. Return a Handle to the underlying env_logger.
pub fn from_env_logger(
    mut builder: env_logger::Builder,
    pass_through: bool,
    hold_time: Option<Duration>,
) -> buff_log::Handle<env_logger::Logger> {
    // Do not have the env_logger set the timestamp. Because the entries are
    // buffered, the timestamp set by the env_logger will correspond to the
    // time at which the entry was dumped, not the time of its origination.
    builder.default_format_timestamp(false);
    buff_log::Logger::new(builder.build(), pass_through, hold_time).init()
}

/// Configure and initialize the logger.
/// If debug is true, log at debug level. Otherwise read log configuration
/// parameters from the environment if RUST_LOG is set. Otherwise, just
/// accept the default configuration.
fn initialize_log(debug: bool) -> buff_log::Handle<env_logger::Logger> {
    let mut builder = Builder::new();
    if debug {
        builder.filter(Some("stratisd"), LevelFilter::Debug);
        builder.filter(Some("libstratis"), LevelFilter::Debug);
        from_env_logger(builder, true, None)
    } else {
        builder.filter_level(LevelFilter::Trace);
        if let Ok(s) = env::var("RUST_LOG") {
            builder.parse(&s);
        }
        from_env_logger(
            builder,
            false,
            Some(Duration::minutes(DEFAULT_LOG_HOLD_MINUTES)),
        )
    }
}

/// To ensure only one instance of stratisd runs at a time, acquire an
/// exclusive lock. Return an error if lock attempt fails.
fn trylock_pid_file() -> StratisResult<File> {
    let mut f = OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .open(STRATISD_PID_PATH)
        .map_err(|err| {
            StratisError::Error(format!(
                "Failed to create or open the stratisd PID file at {}: {}",
                STRATISD_PID_PATH, err
            ))
        })?;
    match flock(f.as_raw_fd(), FlockArg::LockExclusiveNonblock) {
        Ok(_) => {
            f.write_all(getpid().to_string().as_bytes())?;
            Ok(f)
        }
        Err(_) => {
            let mut buf = String::new();
            f.read_to_string(&mut buf)?;

            Err(StratisError::Error(format!(
                "Daemon already running with supposed pid: {}",
                buf
            )))
        }
    }
}

// Conditionally compiled support for a D-Bus interface.
struct MaybeDbusSupport {
    #[cfg(feature = "dbus_enabled")]
    handle: Option<libstratis::dbus_api::DbusConnectionData>,
}

// If D-Bus compiled out, do very little.
#[cfg(not(feature = "dbus_enabled"))]
impl MaybeDbusSupport {
    fn new() -> MaybeDbusSupport {
        MaybeDbusSupport {}
    }

    fn process(
        &mut self,
        _engine: &Rc<RefCell<dyn Engine>>,
        _fds: &mut Vec<libc::pollfd>,
        _dbus_client_index_start: usize,
    ) {
    }

    fn register_pool(&mut self, _pool_uuid: Uuid, _pool: &mut dyn Pool) {}

    fn poll_timeout(&self) -> i32 {
        // Non-DBus timeout is infinite
        -1
    }
}

#[cfg(feature = "dbus_enabled")]
impl MaybeDbusSupport {
    fn new() -> MaybeDbusSupport {
        MaybeDbusSupport { handle: None }
    }

    /// Connect to D-Bus and register pools, if not already connected.
    /// Return the connection, if made or already existing, otherwise, None.
    fn setup_connection(
        &mut self,
        engine: &Rc<RefCell<dyn Engine>>,
    ) -> Option<&mut DbusConnectionData> {
        if self.handle.is_none() {
            match libstratis::dbus_api::DbusConnectionData::connect(Rc::clone(engine)) {
                Err(err) => {
                    warn!("D-Bus API is not available: {}", err);
                }
                Ok(mut handle) => {
                    info!("D-Bus API is available");
                    let event_handler = Box::new(EventHandler::new(Rc::clone(&handle.connection)));
                    get_engine_listener_list_mut().register_listener(event_handler);
                    // Register all the pools with dbus
                    for (_, pool_uuid, pool) in engine.borrow_mut().pools_mut() {
                        handle.register_pool(pool_uuid, pool)
                    }
                    self.handle = Some(handle);
                }
            }
        };
        self.handle.as_mut()
    }

    /// Handle any client dbus requests.
    fn process(
        &mut self,
        engine: &Rc<RefCell<dyn Engine>>,
        fds: &mut Vec<libc::pollfd>,
        dbus_client_index_start: usize,
    ) {
        if let Some(handle) = self.setup_connection(engine) {
            handle.handle(&fds[dbus_client_index_start..]);

            // Refresh list of dbus fds to poll for. This can change as
            // D-Bus clients come and go.
            fds.truncate(dbus_client_index_start);
            fds.extend(
                handle
                    .connection
                    .borrow()
                    .watch_fds()
                    .iter()
                    .map(|w| w.to_pollfd()),
            );
        }
    }

    fn register_pool(&mut self, pool_uuid: Uuid, pool: &mut dyn Pool) {
        if let Some(h) = self.handle.as_mut() {
            h.register_pool(pool_uuid, pool)
        }
    }

    fn poll_timeout(&self) -> i32 {
        // If there is no D-Bus connection set timeout to 1 sec (1000 ms), so
        // that stratisd can periodically attempt to set up a connection.
        // If the connection is up, set the timeout to infinite; there is no
        // need to poll as events will be received.
        self.handle.as_ref().map_or(1000, |_| -1)
    }
}

// A facility for listening for and handling udev events that stratisd
// considers interesting.
struct UdevMonitor<'a> {
    socket: libudev::MonitorSocket<'a>,
}

impl<'a> UdevMonitor<'a> {
    fn create(context: &'a libudev::Context) -> StratisResult<UdevMonitor<'a>> {
        let mut monitor = libudev::Monitor::new(context)?;
        monitor.match_subsystem("block")?;

        Ok(UdevMonitor {
            socket: monitor.listen()?,
        })
    }

    fn as_raw_fd(&mut self) -> RawFd {
        self.socket.as_raw_fd()
    }

    /// Handle udev events.
    /// Check if a pool can be constructed and update engine and D-Bus layer
    /// data structures if so.
    fn handle_events(&mut self, engine: &mut dyn Engine, dbus_support: &mut MaybeDbusSupport) {
        while let Some(event) = self.socket.receive_event() {
            if let Some((pool_uuid, pool)) = engine.handle_event(&event) {
                dbus_support.register_pool(pool_uuid, pool);
            }
        }
    }
}

// Process any pending signals, return true if SIGINT received.
// Return an error if there was an error reading the signal.
fn process_signal(
    sfd: &mut SignalFd,
    buff_log: &buff_log::Handle<env_logger::Logger>,
) -> StratisResult<bool> {
    match sfd.read_signal() {
        // This is an unsafe conversion, but in this context that is
        // mostly harmless. A negative converted value, which is
        // virtually impossible, will not match any of the masked
        // values, and stratisd will panic and exit.
        Ok(Some(sig)) => match sig.ssi_signo as i32 {
            nix::libc::SIGUSR1 => {
                info!(
                    "SIGUSR1 received, dumping {} buffered log entries",
                    buff_log.buffered_count()
                );
                buff_log.dump();
                Ok(false)
            }
            nix::libc::SIGINT => {
                info!("SIGINT received, exiting");
                Ok(true)
            }
            signo => {
                panic!("Caught an impossible signal {:?}", signo);
            }
        },
        // No signals waiting (SFD_NONBLOCK flag is set)
        Ok(None) => Ok(false),

        Err(err) => Err(err.into()),
    }
}

/// Handle blocking the event loop
fn process_poll(poll_timeout: i32, fds: &mut Vec<libc::pollfd>) -> StratisResult<()> {
    let r = unsafe { libc::poll(fds.as_mut_ptr(), fds.len() as libc::c_ulong, poll_timeout) };

    // TODO: refine this behavior.
    // Different behaviors may be indicated, depending on the value of
    // errno when return value is -1.
    if r < 0 {
        return Err(StratisError::Error(format!(
            "poll command failed: number of fds: {}, timeout: {}",
            fds.len(),
            poll_timeout
        )));
    }
    Ok(())
}

/// Set up all sorts of signal and event handling mechanisms.
/// Initialize the engine and keep it running until a signal is received
/// or a fatal error is encountered. Dump log entries on specified signal
/// via buff_log.
fn run(matches: &ArgMatches, buff_log: &buff_log::Handle<env_logger::Logger>) -> StratisResult<()> {
    // Ensure that the debug log is output when we leave this function.
    let _guard = buff_log.to_guard();

    let mut dbus_support = MaybeDbusSupport::new();

    // Setup a udev listener before initializing the engine. A device may
    // appear after the engine has processed the udev db, but before it has
    // completed initialization. Unless the udev event has been recorded, the
    // engine will miss the device.
    // This is especially important since stratisd must run during early boot.
    let context = libudev::Context::new()?;
    let mut udev_monitor = UdevMonitor::create(&context)?;

    let engine: Rc<RefCell<dyn Engine>> = {
        info!("stratis daemon version {} started", VERSION);
        if matches.is_present("sim") {
            info!("Using SimEngine");
            Rc::new(RefCell::new(SimEngine::default()))
        } else {
            info!("Using StratEngine");
            Rc::new(RefCell::new(StratEngine::initialize()?))
        }
    };

    /*
    The file descriptor array indexes are:

    0   == Always udev fd index
    1   == SIGNAL FD index
    2   == engine index if eventable
    2/3 == Start of dbus client file descriptor(s)
            * 2 if engine is not eventable
            * else 3
    */
    const FD_INDEX_UDEV: usize = 0;
    const FD_INDEX_SIGNALFD: usize = 1;
    const FD_INDEX_ENGINE: usize = 2;

    /*
    fds is a Vec of libc::pollfd structs. Ideally, it would be possible
    to use the higher level nix crate to handle polling. If this were possible,
    then the Vec would be one of nix::poll::PollFds and this would be more
    rustic. Unfortunately, the rust D-Bus library requires an explicit file
    descriptor to be passed as an argument to Connection::watch_handle(),
    and the explicit file descriptor can not be extracted from the PollFd
    struct. So, at this time, sticking with libc is less complex than
    converting to using nix, because if using nix, the file descriptor would
    have to be maintained in the Vec as well as the PollFd struct.
    */
    let mut fds = Vec::new();

    fds.push(libc::pollfd {
        fd: udev_monitor.as_raw_fd(),
        revents: 0,
        events: libc::POLLIN,
    });

    // Signals can be queued up on this file descriptor
    let mut sfd = {
        let mut mask = SigSet::empty();
        mask.add(signal::SIGINT);
        mask.add(signal::SIGUSR1);
        mask.thread_block()?;
        SignalFd::with_flags(&mask, SfdFlags::SFD_NONBLOCK)?
    };

    fds.push(libc::pollfd {
        fd: sfd.as_raw_fd(),
        revents: 0,
        events: libc::POLLIN,
    });

    let eventable = engine.borrow().get_eventable();

    if let Some(evt) = eventable {
        fds.push(libc::pollfd {
            fd: evt.get_pollable_fd(),
            revents: 0,
            events: libc::POLLIN,
        });
    };

    let dbus_client_index_start = if eventable.is_some() {
        FD_INDEX_ENGINE + 1
    } else {
        FD_INDEX_ENGINE
    };

    loop {
        if fds[FD_INDEX_UDEV].revents != 0 {
            udev_monitor.handle_events(&mut *engine.borrow_mut(), &mut dbus_support)
        }

        if fds[FD_INDEX_SIGNALFD].revents != 0 {
            match process_signal(&mut sfd, buff_log) {
                Ok(should_exit) => {
                    if should_exit {
                        return Ok(());
                    }
                }
                Err(e) => return Err(e),
            }
        }

        if let Some(evt) = eventable {
            if fds[FD_INDEX_ENGINE].revents != 0 {
                evt.clear_event()?;
                engine.borrow_mut().evented()?;
            }
        }

        dbus_support.process(&engine, &mut fds, dbus_client_index_start);

        process_poll(dbus_support.poll_timeout(), &mut fds)?;
    }
}

fn main() {
    let matches = App::new("stratis")
        .version(VERSION)
        .about("Stratis storage management")
        .arg(
            Arg::with_name("debug")
                .long("debug")
                .help("Print additional output for debugging"),
        )
        .arg(
            Arg::with_name("sim")
                .long("sim")
                .help("Use simulator engine"),
        )
        .get_matches();

    // Using a let-expression here so that the scope of the lock file
    // is the rest of the block.
    let lock_file = trylock_pid_file();

    let result = {
        match lock_file {
            Err(err) => Err(err),
            Ok(_) => {
                let log_handle = initialize_log(matches.is_present("debug"));
                run(&matches, &log_handle)
            }
        }
    };

    if let Err(err) = result {
        print_err(&err);
        exit(1);
    } else {
        exit(0);
    }
}
<|endoftext|>


<commit_before>use arraymap::ArrayMap;
use itertools::Itertools;
use mli::*;
use mli_ndarray::Ndeep;
use ndarray::{s, Array, Array1, Array2, ArrayBase, ArrayView1, Data, OwnedRepr};
use std::marker::PhantomData;

type D1 = ndarray::Ix1;
type D2 = ndarray::Ix2;

fn corners(coordinate: [f32; 2]) -> [[isize; 2]; 4] {
    let c00 = coordinate.map(|f| f.floor() as isize);
    let c11 = coordinate.map(|f| f.ceil() as isize);
    let c01 = [c00[0], c11[1]];
    let c10 = [c11[0], c00[1]];
    [c00, c01, c10, c11]
}

pub struct DefConvData {
    pub offsets: Array2<f32>,
    pub features: Array2<f32>,
}

impl DefConvData {
    fn validate_corners(&self, coordinate: [f32; 2]) -> ([Option<[usize; 2]>; 4], [f32; 2]) {
        // Get the integer versions of four corners.
        let coords = corners(coordinate);

        // Get the relative coordinate.
        let rc = [
            coordinate[0] - coords[0][0] as f32,
            coordinate[1] - coords[0][1] as f32,
        ];

        // Create a closure to detect out-of-bounds.
        let in_bounds = |[y, x]: [isize; 2]| {
            if y < 0
                || x < 0
                || y >= self.features.shape()[0] as isize
                || x >= self.features.shape()[1] as isize
            {
                None
            } else {
                Some([y, x].map(|&n| n as usize))
            }
        };

        // Create a closure to extract the feature with zero-padding.
        let extract = |c| in_bounds(c);

        (coords.map(|&c| extract(c)), rc)
    }

    /// Extracts the corner features and relative coordinate for a bilinear interpolation.
    fn extract_corners(&self, coordinate: [f32; 2]) -> ([f32; 4], [f32; 2]) {
        let (validated_corners, rc) = self.validate_corners(coordinate);
        (
            validated_corners.map(|c| c.as_ref().map(|&c| self.features[c]).unwrap_or(0.0)),
            rc,
        )
    }

    /// Looks up a location in the features using bilinear interpolation and zero-padding.
    fn bilinear(&self, coordinate: [f32; 2]) -> f32 {
        // Extract the corner features and the relative coordinate.
        let (f, rc) = self.extract_corners(coordinate);

        // Perform the y interpolation.
        let fx = [
            (1.0 - rc[0]) * f[0] + rc[0] * f[2],
            (1.0 - rc[0]) * f[1] + rc[0] * f[3],
        ];

        // Perform the x interpolation.
        (1.0 - rc[1]) * fx[0] + rc[1] * fx[1]
    }

    /// Looks up a position gradient from a location in the features using bilinear interpolation and zero-padding.
    fn bilinear_position_gradient(&self, coordinate: [f32; 2]) -> [f32; 2] {
        // Extract the corner features and the relative coordinate.
        let (f, rc) = self.extract_corners(coordinate);

        // Perform the y interpolation to get the x values.
        let fx = [
            (1.0 - rc[0]) * f[0] + rc[0] * f[2],
            (1.0 - rc[0]) * f[1] + rc[0] * f[3],
        ];

        // Perform the x interpolation to get the y values.
        let fy = [
            (1.0 - rc[1]) * f[0] + rc[1] * f[1],
            (1.0 - rc[1]) * f[2] + rc[1] * f[3],
        ];

        // The gradient is the difference between the interpolation values in each dimension.
        [fy[1] - fy[0], fx[1] - fx[0]]
    }
}

#[derive(Clone, Debug)]
pub struct DefConv2 {
    weights: Array1<f32>,
    output_shape: [usize; 2],
}

impl DefConv2 {
    pub fn new(weights: Array1<f32>, output_shape: [usize; 2]) -> Self {
        Self {
            weights,
            output_shape,
        }
    }
}

impl Forward for DefConv2 {
    type Input = DefConvData;
    type Internal = ();
    type Output = Array2<f32>;

    fn forward(&self, input: &Self::Input) -> ((), Self::Output) {
        assert_eq!(self.weights.len(), input.offsets.shape()[0]);
        // Get shapes.
        let outshape = self.output_shape;
        let inshape = input.features.raw_dim();

        // Compute the coordinate multiplier.
        let multipliers = [
            inshape[0] as f32 / outshape[0] as f32,
            inshape[1] as f32 / outshape[1] as f32,
        ];

        // Compute bilinear interpolation for (y, x) pairs.
        (
            (),
            Array::from_shape_vec(
                outshape,
                (0..outshape[0])
                    .cartesian_product(0..outshape[1])
                    .map(|(y, x)| {
                        self.weights
                            .iter()
                            .zip(input.offsets.outer_iter())
                            .map(|(weight, offset)| {
                                input.bilinear([
                                    (y as f32 + 0.5 + offset[0]) * multipliers[0],
                                    (x as f32 + 0.5 + offset[1]) * multipliers[1],
                                ]) * weight
                            })
                            .sum()
                    })
                    .collect(),
            )
            .expect("mli-defconv: unexpected shape vec"),
        )
    }
}

impl Backward for DefConv2 {
    type OutputDelta = Array2<f32>;
    type InputDelta = DefConvData;
    type TrainDelta = Ndeep<OwnedRepr<f32>, D1>;

    fn backward(
        &self,
        input: &Self::Input,
        _: &Self::Internal,
        output_delta: &Self::OutputDelta,
    ) -> (Self::InputDelta, Self::TrainDelta) {
        // Get shapes.
        let outshape = self.output_shape;
        let inshape = input.features.raw_dim();

        // Compute the coordinate multiplier.
        let multipliers = [
            inshape[0] as f32 / outshape[0] as f32,
            inshape[1] as f32 / outshape[1] as f32,
        ];

        // We cannot determine in advance what locations are affected
        let mut feature_deltas: Array2<f32> = Array2::zeros(inshape);

        for (y, x) in (0..outshape[0]).cartesian_product(0..outshape[1]) {
            for (weight, offset) in self.weights.iter().zip(input.offsets.outer_iter()) {
                let sample_coordinate = [
                    (y as f32 + 0.5 + offset[0]) * multipliers[0],
                    (x as f32 + 0.5 + offset[1]) * multipliers[1],
                ];
                let (coords, rc) = input.validate_corners(sample_coordinate);
                for (coord, coeff) in coords.iter().zip(
                    [
                        (1.0 - rc[0]) * (1.0 - rc[1]),
                        (1.0 - rc[0]) * rc[1],
                        rc[0] * (1.0 - rc[1]),
                        rc[0] * rc[1],
                    ]
                    .iter(),
                ) {
                    if let Some(coord) = coord {
                        feature_deltas[*coord] += coeff * weight * output_delta[[y, x]];
                    }
                }
            }
        }

        // Compute bilinear interpolation for (y, x) pairs.
        (
            DefConvData {
                features: feature_deltas,
                offsets: unimplemented!(),
            },
            unimplemented!(),
        )
    }
}
<commit_msg>add offset delta to backprop for defconv<commit_after>use arraymap::ArrayMap;
use itertools::Itertools;
use mli::*;
use mli_ndarray::Ndeep;
use ndarray::{s, Array, Array1, Array2, ArrayBase, ArrayView1, Data, OwnedRepr};
use std::marker::PhantomData;

type D1 = ndarray::Ix1;
type D2 = ndarray::Ix2;

fn corners(coordinate: [f32; 2]) -> [[isize; 2]; 4] {
    let c00 = coordinate.map(|f| f.floor() as isize);
    let c11 = coordinate.map(|f| f.ceil() as isize);
    let c01 = [c00[0], c11[1]];
    let c10 = [c11[0], c00[1]];
    [c00, c01, c10, c11]
}

/// Looks up a position gradient from a location in the features using bilinear interpolation and zero-padding.
fn bilinear_position_gradient(f: [f32; 4], rc: [f32; 2]) -> [f32; 2] {
    // Perform the y interpolation to get the x values.
    let fx = [
        (1.0 - rc[0]) * f[0] + rc[0] * f[2],
        (1.0 - rc[0]) * f[1] + rc[0] * f[3],
    ];

    // Perform the x interpolation to get the y values.
    let fy = [
        (1.0 - rc[1]) * f[0] + rc[1] * f[1],
        (1.0 - rc[1]) * f[2] + rc[1] * f[3],
    ];

    // The gradient is the difference between the interpolation values in each dimension.
    [fy[1] - fy[0], fx[1] - fx[0]]
}

pub struct DefConvData {
    pub offsets: Array2<f32>,
    pub features: Array2<f32>,
}

impl DefConvData {
    fn validate_corners(&self, coordinate: [f32; 2]) -> ([Option<[usize; 2]>; 4], [f32; 2]) {
        // Get the integer versions of four corners.
        let coords = corners(coordinate);

        // Get the relative coordinate.
        let rc = [
            coordinate[0] - coords[0][0] as f32,
            coordinate[1] - coords[0][1] as f32,
        ];

        // Create a closure to detect out-of-bounds.
        let in_bounds = |[y, x]: [isize; 2]| {
            if y < 0
                || x < 0
                || y >= self.features.shape()[0] as isize
                || x >= self.features.shape()[1] as isize
            {
                None
            } else {
                Some([y, x].map(|&n| n as usize))
            }
        };

        // Create a closure to extract the feature with zero-padding.
        let extract = |c| in_bounds(c);

        (coords.map(|&c| extract(c)), rc)
    }

    /// Extracts the corner features and relative coordinate for a bilinear interpolation.
    fn extract_corners(&self, validated_corners: [Option<[usize; 2]>; 4]) -> [f32; 4] {
        validated_corners.map(|c| c.as_ref().map(|&c| self.features[c]).unwrap_or(0.0))
    }

    /// Looks up a location in the features using bilinear interpolation and zero-padding.
    fn bilinear(&self, coordinate: [f32; 2]) -> f32 {
        // Extract the corner features and the relative coordinate.
        let (validated_corners, rc) = self.validate_corners(coordinate);
        let f = self.extract_corners(validated_corners);

        // Perform the y interpolation.
        let fx = [
            (1.0 - rc[0]) * f[0] + rc[0] * f[2],
            (1.0 - rc[0]) * f[1] + rc[0] * f[3],
        ];

        // Perform the x interpolation.
        (1.0 - rc[1]) * fx[0] + rc[1] * fx[1]
    }
}

#[derive(Clone, Debug)]
pub struct DefConv2 {
    weights: Array1<f32>,
    output_shape: [usize; 2],
}

impl DefConv2 {
    pub fn new(weights: Array1<f32>, output_shape: [usize; 2]) -> Self {
        Self {
            weights,
            output_shape,
        }
    }
}

impl Forward for DefConv2 {
    type Input = DefConvData;
    type Internal = ();
    type Output = Array2<f32>;

    fn forward(&self, input: &Self::Input) -> ((), Self::Output) {
        assert_eq!(self.weights.len(), input.offsets.shape()[0]);
        // Get shapes.
        let outshape = self.output_shape;
        let inshape = input.features.raw_dim();

        // Compute the coordinate multiplier.
        let multipliers = [
            inshape[0] as f32 / outshape[0] as f32,
            inshape[1] as f32 / outshape[1] as f32,
        ];

        // Compute bilinear interpolation for (y, x) pairs.
        (
            (),
            Array::from_shape_vec(
                outshape,
                (0..outshape[0])
                    .cartesian_product(0..outshape[1])
                    .map(|(y, x)| {
                        self.weights
                            .iter()
                            .zip(input.offsets.outer_iter())
                            .map(|(weight, offset)| {
                                input.bilinear([
                                    (y as f32 + 0.5 + offset[0]) * multipliers[0],
                                    (x as f32 + 0.5 + offset[1]) * multipliers[1],
                                ]) * weight
                            })
                            .sum()
                    })
                    .collect(),
            )
            .expect("mli-defconv: unexpected shape vec"),
        )
    }
}

impl Backward for DefConv2 {
    type OutputDelta = Array2<f32>;
    type InputDelta = DefConvData;
    type TrainDelta = Ndeep<OwnedRepr<f32>, D1>;

    fn backward(
        &self,
        input: &Self::Input,
        _: &Self::Internal,
        output_delta: &Self::OutputDelta,
    ) -> (Self::InputDelta, Self::TrainDelta) {
        // Get shapes.
        let outshape = self.output_shape;
        let inshape = input.features.raw_dim();

        // Compute the coordinate multiplier.
        let multipliers = [
            inshape[0] as f32 / outshape[0] as f32,
            inshape[1] as f32 / outshape[1] as f32,
        ];

        // We cannot determine in advance what locations are affected
        let mut feature_deltas: Array2<f32> = Array2::zeros(inshape);
        let mut offset_deltas: Array2<f32> = Array2::zeros(input.offsets.raw_dim());

        for (y, x) in (0..outshape[0]).cartesian_product(0..outshape[1]) {
            let output_delta = output_delta[[y, x]];
            for (ix, (weight, offset)) in self
                .weights
                .iter()
                .zip(input.offsets.outer_iter())
                .enumerate()
            {
                // Compute the original sample coordinate.
                let sample_coordinate = [
                    (y as f32 + 0.5 + offset[0]) * multipliers[0],
                    (x as f32 + 0.5 + offset[1]) * multipliers[1],
                ];

                // Find the corners if they are within the input.
                let (validated_corners, rc) = input.validate_corners(sample_coordinate);

                // Loop over all the corners and the amount each contributed.
                for (coord, coeff) in validated_corners.iter().zip(
                    [
                        (1.0 - rc[0]) * (1.0 - rc[1]),
                        (1.0 - rc[0]) * rc[1],
                        rc[0] * (1.0 - rc[1]),
                        rc[0] * rc[1],
                    ]
                    .iter(),
                ) {
                    // If the coordinate was in bounds of the input tensor.
                    if let Some(coord) = coord {
                        // Add the contribution.
                        feature_deltas[*coord] += coeff * weight * output_delta;
                    }
                }

                // Compute the position gradients.
                let f = input.extract_corners(validated_corners);
                let position_gradient = bilinear_position_gradient(f, rc);
                // Compute and add the offset gradients based on the position gradients and chain rule.
                // The multipliers affect the gradient of the position because they multiply the sampling locations.
                // The greater the weight the greater the effect on sampling location because the output is multiplied by the weight.
                offset_deltas[[ix, 0]] +=
                    multipliers[0] * position_gradient[0] * weight * output_delta;
                offset_deltas[[ix, 1]] +=
                    multipliers[1] * position_gradient[1] * weight * output_delta;
            }
        }

        // Compute bilinear interpolation for (y, x) pairs.
        (
            DefConvData {
                features: feature_deltas,
                offsets: offset_deltas,
            },
            unimplemented!(),
        )
    }
}
<|endoftext|>


<commit_before>//! Types for working with [`File`].
//!
//! [`File`]: file/struct.File.html

mod create;
mod metadata;
mod open;
mod open_options;
mod seek;

pub use self::create::CreateFuture;
pub use self::metadata::MetadataFuture;
pub use self::open::OpenFuture;
pub use self::open_options::OpenOptions;
pub use self::seek::SeekFuture;

use tokio_io::{AsyncRead, AsyncWrite};

use futures::Poll;

use std::fs::{File as StdFile, Metadata, Permissions};
use std::io::{self, Read, Write, Seek};
use std::path::Path;

/// A reference to an open file on the filesystem.
///
/// This is a specialized version of [`std::fs::File`][std] for usage from the
/// Tokio runtime.
///
/// An instance of a `File` can be read and/or written depending on what options
/// it was opened with. Files also implement Seek to alter the logical cursor
/// that the file contains internally.
///
/// Files are automatically closed when they go out of scope.
///
/// [std]: https://doc.rust-lang.org/std/fs/struct.File.html
#[derive(Debug)]
pub struct File {
    std: Option<StdFile>,
}

impl File {
    /// Attempts to open a file in read-only mode.
    ///
    /// See [`OpenOptions`] for more details.
    ///
    /// [`OpenOptions`]: struct.OpenOptions.html
    ///
    /// # Errors
    ///
    /// `OpenFuture` results in an error if called from outside of the Tokio
    /// runtime or if the underlying [`open`] call results in an error.
    ///
    /// [`open`]: https://doc.rust-lang.org/std/fs/struct.File.html#method.open
    pub fn open
(path: P) -> OpenFuture

    where P: AsRef<Path> + Send + 'static,
    {
        OpenOptions::new().read(true).open(path)
    }

    /// Opens a file in write-only mode.
    ///
    /// This function will create a file if it does not exist, and will truncate
    /// it if it does.
    ///
    /// See [`OpenOptions`] for more details.
    ///
    /// [`OpenOptions`]: struct.OpenOptions.html
    ///
    /// # Errors
    ///
    /// `CreateFuture` results in an error if called from outside of the Tokio
    /// runtime or if the underlying [`create`] call results in an error.
    ///
    /// [`create`]: https://doc.rust-lang.org/std/fs/struct.File.html#method.create
    pub fn create
(path: P) -> CreateFuture

    where P: AsRef<Path> + Send + 'static,
    {
        CreateFuture::new(path)
    }

    /// Convert a [`std::fs::File`][std] to a `tokio_fs::File`.
    ///
    /// [std]: https://doc.rust-lang.org/std/fs/struct.File.html
    pub fn from_std(std: StdFile) -> File {
        File { std: Some(std) }
    }

    /// Seek to an offset, in bytes, in a stream.
    ///
    /// A seek beyond the end of a stream is allowed, but implementation
    /// defined.
    ///
    /// If the seek operation completed successfully, this method returns the
    /// new position from the start of the stream. That position can be used
    /// later with `SeekFrom::Start`.
    ///
    /// # Errors
    ///
    /// Seeking to a negative offset is considered an error.
    pub fn poll_seek(&mut self, pos: io::SeekFrom) -> Poll<u64, io::Error> {
        ::blocking_io(|| self.std().seek(pos))
    }

    /// Seek to an offset, in bytes, in a stream.
    ///
    /// Similar to `poll_seek`, but returning a `Future`.
    ///
    /// This method consumes the `File` and returns it back when the future
    /// completes.
    pub fn seek(self, pos: io::SeekFrom) -> SeekFuture {
        SeekFuture::new(self, pos)
    }

    /// Attempts to sync all OS-internal metadata to disk.
    ///
    /// This function will attempt to ensure that all in-core data reaches the
    /// filesystem before returning.
    pub fn poll_sync_all(&mut self) -> Poll<(), io::Error> {
        ::blocking_io(|| self.std().sync_all())
    }

    /// This function is similar to `poll_sync_all`, except that it may not
    /// synchronize file metadata to the filesystem.
    ///
    /// This is intended for use cases that must synchronize content, but don't
    /// need the metadata on disk. The goal of this method is to reduce disk
    /// operations.
    ///
    /// Note that some platforms may simply implement this in terms of `poll_sync_all`.
    pub fn poll_sync_data(&mut self) -> Poll<(), io::Error> {
        ::blocking_io(|| self.std().sync_data())
    }

    /// Truncates or extends the underlying file, updating the size of this file to become size.
    ///
    /// If the size is less than the current file's size, then the file will be
    /// shrunk. If it is greater than the current file's size, then the file
    /// will be extended to size and have all of the intermediate data filled in
    /// with 0s.
    ///
    /// # Errors
    ///
    /// This function will return an error if the file is not opened for
    /// writing.
    pub fn poll_set_len(&mut self, size: u64) -> Poll<(), io::Error> {
        ::blocking_io(|| self.std().set_len(size))
    }

    /// Queries metadata about the underlying file.
    pub fn metadata(self) -> MetadataFuture {
        MetadataFuture::new(self)
    }

    /// Queries metadata about the underlying file.
    pub fn poll_metadata(&mut self) -> Poll<Metadata, io::Error> {
        ::blocking_io(|| self.std().metadata())
    }

    /// Create a new `File` instance that shares the same underlying file handle
    /// as the existing `File` instance. Reads, writes, and seeks will affect both
    /// File instances simultaneously.
    pub fn poll_try_clone(&mut self) -> Poll<File, io::Error> {
        ::blocking_io(|| {
            let std = self.std().try_clone()?;
            Ok(File::from_std(std))
        })
    }

    /// Changes the permissions on the underlying file.
    ///
    /// # Platform-specific behavior
    ///
    /// This function currently corresponds to the `fchmod` function on Unix and
    /// the `SetFileInformationByHandle` function on Windows. Note that, this
    /// [may change in the future][changes].
    ///
    /// [changes]: https://doc.rust-lang.org/std/io/index.html#platform-specific-behavior
    ///
    /// # Errors
    ///
    /// This function will return an error if the user lacks permission change
    /// attributes on the underlying file. It may also return an error in other
    /// os-specific unspecified cases.
    pub fn poll_set_permissions(&mut self, perm: Permissions) -> Poll<(), io::Error> {
        ::blocking_io(|| self.std().set_permissions(perm))
    }

    /// Destructures the `tokio_fs::File` into a [`std::fs::File`][std].
    ///
    /// # Panics
    ///
    /// This function will panic if `shutdown` has been called.
    ///
    /// [std]: https://doc.rust-lang.org/std/fs/struct.File.html
    pub fn into_std(mut self) -> StdFile {
        self.std.take().expect("`File` instance already shutdown")
    }

    fn std(&mut self) -> &mut StdFile {
        self.std.as_mut().expect("`File` instance already shutdown")
    }
}

impl Read for File {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        ::would_block(|| self.std().read(buf))
    }
}

impl AsyncRead for File {
    unsafe fn prepare_uninitialized_buffer(&self, _: &mut [u8]) -> bool {
        false
    }
}

impl Write for File {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        ::would_block(|| self.std().write(buf))
    }

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

impl AsyncWrite for File {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        ::blocking_io(|| {
            self.std = None;
            Ok(())
        })
    }
}

impl Drop for File {
    fn drop(&mut self) {
        if let Some(_std) = self.std.take() {
            // This is probably fine as closing a file *shouldn't* be a blocking
            // operation. That said, ideally `shutdown` is called first.
        }
    }
}
<commit_msg>fs: added usage examples/doctests to File (#786)<commit_after>//! Types for working with [`File`].
//!
//! [`File`]: file/struct.File.html

mod create;
mod metadata;
mod open;
mod open_options;
mod seek;

pub use self::create::CreateFuture;
pub use self::metadata::MetadataFuture;
pub use self::open::OpenFuture;
pub use self::open_options::OpenOptions;
pub use self::seek::SeekFuture;

use tokio_io::{AsyncRead, AsyncWrite};

use futures::Poll;

use std::fs::{File as StdFile, Metadata, Permissions};
use std::io::{self, Read, Write, Seek};
use std::path::Path;

/// A reference to an open file on the filesystem.
///
/// This is a specialized version of [`std::fs::File`][std] for usage from the
/// Tokio runtime.
///
/// An instance of a `File` can be read and/or written depending on what options
/// it was opened with. Files also implement Seek to alter the logical cursor
/// that the file contains internally.
///
/// Files are automatically closed when they go out of scope.
///
/// [std]: https://doc.rust-lang.org/std/fs/struct.File.html
///
/// # Examples
///
/// Create a new file and asynchronously write bytes to it:
///
/// ```no_run
/// extern crate tokio;
///
/// use tokio::prelude::{AsyncWrite, Future};
///
/// fn main() {
///     let task = tokio::fs::File::create("foo.txt")
///         .and_then(|mut file| file.poll_write(b"hello, world!"))
///         .map(|res| {
///             println!("{:?}", res);
///         }).map_err(|err| eprintln!("IO error: {:?}", err));
///
///     tokio::run(task);
/// }
/// ```
///
/// Read the contents of a file into a buffer
///
/// ```no_run
/// extern crate tokio;
///
/// use tokio::prelude::{AsyncRead, Future};
///
/// fn main() {
///     let task = tokio::fs::File::open("foo.txt")
///         .and_then(|mut file| {
///             let mut contents = vec![];
///             file.read_buf(&mut contents)
///                 .map(|res| {
///                     println!("{:?}", res);
///                 })
///         }).map_err(|err| eprintln!("IO error: {:?}", err));
///     tokio::run(task);
/// }
/// ```
#[derive(Debug)]
pub struct File {
    std: Option<StdFile>,
}

impl File {
    /// Attempts to open a file in read-only mode.
    ///
    /// See [`OpenOptions`] for more details.
    ///
    /// [`OpenOptions`]: struct.OpenOptions.html
    ///
    /// # Errors
    ///
    /// `OpenFuture` results in an error if called from outside of the Tokio
    /// runtime or if the underlying [`open`] call results in an error.
    ///
    /// [`open`]: https://doc.rust-lang.org/std/fs/struct.File.html#method.open
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate tokio;
    /// use tokio::prelude::Future;
    /// fn main() {
    ///     let task = tokio::fs::File::open("foo.txt").and_then(|file| {
    ///         // do something with the file ...
    ///         file.metadata().map(|md| println!("{:?}", md))
    ///     }).map_err(|e| {
    ///         // handle errors
    ///         eprintln!("IO error: {:?}", e);
    ///     });
    ///     tokio::run(task);
    /// }
    /// ```
    pub fn open
(path: P) -> OpenFuture

    where P: AsRef<Path> + Send + 'static,
    {
        OpenOptions::new().read(true).open(path)
    }

    /// Opens a file in write-only mode.
    ///
    /// This function will create a file if it does not exist, and will truncate
    /// it if it does.
    ///
    /// See [`OpenOptions`] for more details.
    ///
    /// [`OpenOptions`]: struct.OpenOptions.html
    ///
    /// # Errors
    ///
    /// `CreateFuture` results in an error if called from outside of the Tokio
    /// runtime or if the underlying [`create`] call results in an error.
    ///
    /// [`create`]: https://doc.rust-lang.org/std/fs/struct.File.html#method.create
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate tokio;
    /// use tokio::prelude::Future;
    /// fn main() {
    ///     let task = tokio::fs::File::create("foo.txt")
    ///         .and_then(|file| {
    ///             // do something with the created file ...
    ///             file.metadata().map(|md| println!("{:?}", md))
    ///         }).map_err(|e| {
    ///             // handle errors
    ///             eprintln!("IO error: {:?}", e);
    ///     });
    ///     tokio::run(task);
    /// }
    /// ```
    pub fn create
(path: P) -> CreateFuture

    where P: AsRef<Path> + Send + 'static,
    {
        CreateFuture::new(path)
    }

    /// Convert a [`std::fs::File`][std] to a [`tokio_fs::File`][file].
    ///
    /// [std]: https://doc.rust-lang.org/std/fs/struct.File.html
    /// [file]: struct.File.html
    ///
    /// Examples
    /// ```no_run
    /// # extern crate tokio;
    /// use std::fs::File;
    ///
    /// fn main() {
    ///     let std_file = File::open("foo.txt").unwrap();
    ///     let file = tokio::fs::File::from_std(std_file);
    /// }
    /// ```
    pub fn from_std(std: StdFile) -> File {
        File { std: Some(std) }
    }

    /// Seek to an offset, in bytes, in a stream.
    ///
    /// A seek beyond the end of a stream is allowed, but implementation
    /// defined.
    ///
    /// If the seek operation completed successfully, this method returns the
    /// new position from the start of the stream. That position can be used
    /// later with `SeekFrom::Start`.
    ///
    /// # Errors
    ///
    /// Seeking to a negative offset is considered an error.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate tokio;
    /// use tokio::prelude::Future;
    /// use std::io::SeekFrom;
    ///
    /// fn main() {
    ///     let task = tokio::fs::File::open("foo.txt")
    ///         // move cursor 6 bytes from the start of the file
    ///         .and_then(|mut file| file.poll_seek(SeekFrom::Start(6)))
    ///         .map(|res| {
    ///             println!("{:?}", res);
    ///         }).map_err(|err| eprintln!("IO error: {:?}", err));
    ///
    ///     tokio::run(task);
    /// }
    /// ```
    pub fn poll_seek(&mut self, pos: io::SeekFrom) -> Poll<u64, io::Error> {
        ::blocking_io(|| self.std().seek(pos))
    }

    /// Seek to an offset, in bytes, in a stream.
    ///
    /// Similar to `poll_seek`, but returning a `Future`.
    ///
    /// This method consumes the `File` and returns it back when the future
    /// completes.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate tokio;
    /// use tokio::prelude::Future;
    /// use std::io::SeekFrom;
    ///
    /// fn main() {
    ///     let task = tokio::fs::File::create("foo.txt")
    ///         .and_then(|file| file.seek(SeekFrom::Start(6)))
    ///         .map(|file| {
    ///             // handle returned file ..
    ///             # println!("{:?}", file);
    ///         }).map_err(|err| eprintln!("IO error: {:?}", err));
    ///
    ///     tokio::run(task);
    /// }
    /// ```
    pub fn seek(self, pos: io::SeekFrom) -> SeekFuture {
        SeekFuture::new(self, pos)
    }

    /// Attempts to sync all OS-internal metadata to disk.
    ///
    /// This function will attempt to ensure that all in-core data reaches the
    /// filesystem before returning.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate tokio;
    /// use tokio::prelude::{AsyncWrite, Future};
    ///
    /// fn main() {
    ///     let task = tokio::fs::File::create("foo.txt")
    ///         .and_then(|mut file| {
    ///             file.poll_write(b"hello, world!")?;
    ///             file.poll_sync_all()
    ///         })
    ///         .map(|res| {
    ///             // handle returned result ..
    ///             # println!("{:?}", res);
    ///         }).map_err(|err| eprintln!("IO error: {:?}", err));
    ///
    ///     tokio::run(task);
    /// }
    /// ```
    pub fn poll_sync_all(&mut self) -> Poll<(), io::Error> {
        ::blocking_io(|| self.std().sync_all())
    }

    /// This function is similar to `poll_sync_all`, except that it may not
    /// synchronize file metadata to the filesystem.
    ///
    /// This is intended for use cases that must synchronize content, but don't
    /// need the metadata on disk. The goal of this method is to reduce disk
    /// operations.
    ///
    /// Note that some platforms may simply implement this in terms of `poll_sync_all`.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate tokio;
    /// use tokio::prelude::{AsyncWrite, Future};
    ///
    /// fn main() {
    ///     let task = tokio::fs::File::create("foo.txt")
    ///         .and_then(|mut file| {
    ///             file.poll_write(b"hello, world!")?;
    ///             file.poll_sync_data()
    ///         })
    ///         .map(|res| {
    ///             // handle returned result ..
    ///             # println!("{:?}", res);
    ///         }).map_err(|err| eprintln!("IO error: {:?}", err));
    ///
    ///     tokio::run(task);
    /// }
    /// ```
    pub fn poll_sync_data(&mut self) -> Poll<(), io::Error> {
        ::blocking_io(|| self.std().sync_data())
    }

    /// Truncates or extends the underlying file, updating the size of this file to become size.
    ///
    /// If the size is less than the current file's size, then the file will be
    /// shrunk. If it is greater than the current file's size, then the file
    /// will be extended to size and have all of the intermediate data filled in
    /// with 0s.
    ///
    /// # Errors
    ///
    /// This function will return an error if the file is not opened for
    /// writing.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate tokio;
    /// use tokio::prelude::Future;
    ///
    /// fn main() {
    ///     let task = tokio::fs::File::create("foo.txt")
    ///         .and_then(|mut file| {
    ///             file.poll_set_len(10)
    ///         })
    ///         .map(|res| {
    ///             // handle returned result ..
    ///             # println!("{:?}", res);
    ///         }).map_err(|err| eprintln!("IO error: {:?}", err));
    ///
    ///     tokio::run(task);
    /// }
    /// ```
    pub fn poll_set_len(&mut self, size: u64) -> Poll<(), io::Error> {
        ::blocking_io(|| self.std().set_len(size))
    }

    /// Queries metadata about the underlying file.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate tokio;
    /// use tokio::prelude::Future;
    ///
    /// fn main() {
    ///     let task = tokio::fs::File::create("foo.txt")
    ///         .and_then(|file| file.metadata())
    ///         .map(|metadata| {
    ///             println!("{:?}", metadata);
    ///         }).map_err(|err| eprintln!("IO error: {:?}", err));
    ///
    ///     tokio::run(task);
    /// }
    /// ```
    pub fn metadata(self) -> MetadataFuture {
        MetadataFuture::new(self)
    }

    /// Queries metadata about the underlying file.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate tokio;
    /// use tokio::prelude::Future;
    ///
    /// fn main() {
    ///     let task = tokio::fs::File::create("foo.txt")
    ///         .and_then(|mut file| file.poll_metadata())
    ///         .map(|metadata| {
    ///             // metadata is of type Async::Ready<Metadata>
    ///             println!("{:?}", metadata);
    ///         }).map_err(|err| eprintln!("IO error: {:?}", err));
    ///
    ///     tokio::run(task);
    /// }
    /// ```
    pub fn poll_metadata(&mut self) -> Poll<Metadata, io::Error> {
        ::blocking_io(|| self.std().metadata())
    }

    /// Create a new `File` instance that shares the same underlying file handle
    /// as the existing `File` instance. Reads, writes, and seeks will affect both
    /// File instances simultaneously.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate tokio;
    /// use tokio::prelude::Future;
    ///
    /// fn main() {
    ///     let task = tokio::fs::File::create("foo.txt")
    ///         .and_then(|mut file| file.poll_try_clone())
    ///         .map(|clone| {
    ///             // do something with the clone
    ///             # println!("{:?}", clone);
    ///         }).map_err(|err| eprintln!("IO error: {:?}", err));
    ///
    ///     tokio::run(task);
    /// }
    /// ```
    pub fn poll_try_clone(&mut self) -> Poll<File, io::Error> {
        ::blocking_io(|| {
            let std = self.std().try_clone()?;
            Ok(File::from_std(std))
        })
    }

    /// Changes the permissions on the underlying file.
    ///
    /// # Platform-specific behavior
    ///
    /// This function currently corresponds to the `fchmod` function on Unix and
    /// the `SetFileInformationByHandle` function on Windows. Note that, this
    /// [may change in the future][changes].
    ///
    /// [changes]: https://doc.rust-lang.org/std/io/index.html#platform-specific-behavior
    ///
    /// # Errors
    ///
    /// This function will return an error if the user lacks permission change
    /// attributes on the underlying file. It may also return an error in other
    /// os-specific unspecified cases.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate tokio;
    /// use tokio::prelude::Future;
    ///
    /// fn main() {
    ///     let task = tokio::fs::File::create("foo.txt")
    ///         .and_then(|file| file.metadata())
    ///         .map(|(mut file, metadata)| {
    ///             let mut perms = metadata.permissions();
    ///             perms.set_readonly(true);
    ///             match file.poll_set_permissions(perms) {
    ///                 Err(e) => eprintln!("{}", e),
    ///                 _ => println!("permissions set!"),
    ///             }
    ///         }).map_err(|err| eprintln!("IO error: {:?}", err));
    ///
    ///     tokio::run(task);
    /// }
    /// ```
    pub fn poll_set_permissions(&mut self, perm: Permissions) -> Poll<(), io::Error> {
        ::blocking_io(|| self.std().set_permissions(perm))
    }

    /// Destructures the `tokio_fs::File` into a [`std::fs::File`][std].
    ///
    /// # Panics
    ///
    /// This function will panic if `shutdown` has been called.
    ///
    /// [std]: https://doc.rust-lang.org/std/fs/struct.File.html
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate tokio;
    /// use tokio::prelude::Future;
    ///
    /// fn main() {
    ///     let task = tokio::fs::File::create("foo.txt")
    ///         .map(|file| {
    ///             let std_file = file.into_std();
    ///             // do something with the std::fs::File
    ///             # println!("{:?}", std_file);
    ///         }).map_err(|err| eprintln!("IO error: {:?}", err));
    ///
    ///     tokio::run(task);
    /// }
    /// ```
    pub fn into_std(mut self) -> StdFile {
        self.std.take().expect("`File` instance already shutdown")
    }

    fn std(&mut self) -> &mut StdFile {
        self.std.as_mut().expect("`File` instance already shutdown")
    }
}

impl Read for File {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        ::would_block(|| self.std().read(buf))
    }
}

impl AsyncRead for File {
    unsafe fn prepare_uninitialized_buffer(&self, _: &mut [u8]) -> bool {
        false
    }
}

impl Write for File {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        ::would_block(|| self.std().write(buf))
    }

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

impl AsyncWrite for File {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        ::blocking_io(|| {
            self.std = None;
            Ok(())
        })
    }
}

impl Drop for File {
    fn drop(&mut self) {
        if let Some(_std) = self.std.take() {
            // This is probably fine as closing a file *shouldn't* be a blocking
            // operation. That said, ideally `shutdown` is called first.
        }
    }
}
<|endoftext|>


<commit_before>use num::Integer;
use euclid::Point2D;
use hex::Dir6;

pub trait FovValue: Eq + Clone {
    /// Construct a new FovValue based on previous one and a new point in the fov.
    fn advance(&self, offset: Point2D<i32>) -> Option<Self>;

    /// Optional method for showing acute corners of fake isometric rooms in fov.
    fn is_fake_isometric_wall(&self, offset: Point2D<i32>) -> bool {
        let _ = offset;
        false
    }
}

/// Field of view iterator for a hexagonal map.
///
/// Takes a function that maps cells into user-specified values and indicates when traversal should
/// stop.
pub struct HexFov<T> {
    stack: Vec<Sector<T>>,
    /// Extra values generated by special cases.
    side_channel: Vec<(Point2D<i32>, T)>,
}

impl<T: FovValue> HexFov<T> {
    pub fn new(init: T) -> HexFov<T> {
        // We could run f for (0, 0) here, but the traditional way for the FOV to work is to only
        // consider your surroundings, not the origin site itself.
        HexFov {
            stack: vec![Sector {
                            begin: PolarPoint::new(0.0, 1),
                            pt: PolarPoint::new(0.0, 1),
                            end: PolarPoint::new(6.0, 1),
                            prev_value: init.clone(),
                            group_value: None,
                        }],
            // The FOV algorithm will not generate the origin point, so we use
            // the side channel to explicitly add it in the beginning.
            side_channel: vec![(Point2D::new(0, 0), init)],
        }
    }
}

impl<T: FovValue> Iterator for HexFov<T> {
    type Item = (Point2D<i32>, T);
    fn next(&mut self) -> Option<(Point2D<i32>, T)> {
        // Empty the side channel before proceeding to the algorithm proper.
        if let Some(ret) = self.side_channel.pop() {
            return Some(ret);
        }

        // Start processing the next sector in the stack.
        if let Some(mut current) = self.stack.pop() {
            // Generate the user value from the previous value in the stack item and the working
            // position in the current sector.
            let next_value = current.prev_value.advance(current.pt.to_v2());

            // Proceed along the sector if we're not at its end yet.
            if current.pt.is_below(current.end) {
                let pos = current.pt.to_v2();

                match current.group_value {
                    None => {
                        // Beginning of group, value isn't set.
                        current.group_value = next_value.clone();
                    }
                    ref old_value if old_value != &next_value => {
                        // Value changed, branch out.

                        // Add the rest of this sector with the new value.
                        self.stack.push(Sector {
                            begin: current.pt,
                            pt: current.pt,
                            end: current.end,
                            prev_value: current.prev_value,
                            group_value: next_value.clone(),
                        });

                        // Branch further on the arc processed so far.
                        if let &Some(ref old_value) = old_value {
                            self.stack.push(Sector {
                                begin: current.begin.further(),
                                pt: current.begin.further(),
                                end: current.pt.further(),
                                prev_value: old_value.clone(),
                                group_value: None,
                            });
                        }
                        return self.next();
                    }
                    // Otherwise don't modify the stack.
                    _ => {}
                }

                // Current value and group value are assumed to be the same from here on.
                assert!(current.group_value == next_value);

                // Hack for making acute corner tiles of fake-isometric rooms visible.

                // We're moving along a vertical line on the hex circle, so there are side
                // points to check.
                if let Some(side_pos) = current.pt.side_point() {

                    let next = current.pt.next();
                    // If the next cell is within the current span and the current cell is
                    // wallform,
                    if next.is_below(current.end) &&
                       current.prev_value.is_fake_isometric_wall(current.pt.to_v2()) {
                        // and if the next cell is visible,
                        if let Some(next_value) = current.prev_value.advance(next.to_v2()) {
                            // and if the current and the next cell are in the same value group,
                            // both the next cell and the third corner point cell are
                            // wallforms, and the side point would not be otherwise
                            // visible:
                            if next_value == current.prev_value &&
                               next_value.is_fake_isometric_wall(next.to_v2()) &&
                               current.prev_value.advance(side_pos).is_none() &&
                               current.prev_value.is_fake_isometric_wall(side_pos) {
                                // Add the side point to the side channel.
                                self.side_channel.push((side_pos, current.prev_value.clone()));
                            }
                        }
                    }
                }

                // Proceed along the current sector.
                current.pt = current.pt.next();
                self.stack.push(current);
                if let Some(next_value) = next_value {
                    return Some((pos, next_value));
                } else {
                    return self.next();
                }
            } else {
                // Hit the end of the sector, branch ahead.
                if let Some(value) = current.group_value {
                    self.stack.push(Sector {
                        begin: current.begin.further(),
                        pt: current.begin.further(),
                        end: current.end.further(),
                        prev_value: value,
                        group_value: None,
                    });
                }

                self.next()
            }
        } else {
            None
        }
    }
}

struct Sector<T> {
    /// Start point of current sector.
    begin: PolarPoint,
    /// Point currently being processed.
    pt: PolarPoint,
    /// End point of current sector.
    end: PolarPoint,
    /// The user value from previous iteration.
    prev_value: T,
    /// The user value for this group.
    group_value: Option<T>,
}

/// Points on a hex circle expressed in polar coordinates.
#[derive(Copy, Clone, PartialEq)]
struct PolarPoint {
    pos: f32,
    radius: u32,
}

impl PolarPoint {
    pub fn new(pos: f32, radius: u32) -> PolarPoint {
        PolarPoint {
            pos: pos,
            radius: radius,
        }
    }
    /// Index of the discrete hex cell along the circle that corresponds to this point.
    fn winding_index(self) -> i32 { (self.pos + 0.5).floor() as i32 }

    pub fn is_below(self, other: PolarPoint) -> bool { self.winding_index() < other.end_index() }

    fn end_index(self) -> i32 { (self.pos + 0.5).ceil() as i32 }

    pub fn to_v2(self) -> Point2D<i32> {
        if self.radius == 0 {
            return Point2D::new(0, 0);
        }
        let index = self.winding_index();
        let sector = index.mod_floor(&(self.radius as i32 * 6)) / self.radius as i32;
        let offset = index.mod_floor(&(self.radius as i32));

        let rod = Dir6::from_int(sector).to_v2();
        let tangent = Dir6::from_int((sector + 2) % 6).to_v2();

        rod * (self.radius as i32) + tangent * offset
    }

    /// If this point and the next point are adjacent vertically (along the xy
    /// axis), return the point outside of the circle between the two points.
    ///
    /// This is a helper function for the FOV special case where acute corners
    /// of fake isometric rooms are marked visible even though strict hex FOV
    /// logic would keep them unseen.
    pub fn side_point(self) -> Option<Point2D<i32>> {
        let next = self.next();
        let a = self.to_v2();
        let b = next.to_v2();

        if b.x == a.x + 1 && b.y == a.y + 1 {
            // Going down the right rim.
            Some(Point2D::new(a.x + 1, a.y))
        } else if b.x == a.x - 1 && b.y == a.y - 1 {
            // Going up the left rim.
            Some(Point2D::new(a.x - 1, a.y))
        } else {
            None
        }
    }

    /// The point corresponding to this one on the hex circle with radius +1.
    pub fn further(self) -> PolarPoint {
        PolarPoint::new(self.pos * (self.radius + 1) as f32 / self.radius as f32,
                        self.radius + 1)
    }

    /// The point next to this one along the hex circle.
    pub fn next(self) -> PolarPoint { PolarPoint::new((self.pos + 0.5).floor() + 0.5, self.radius) }
}

#[cfg(test)]
mod test {
    use std::collections::HashMap;
    use std::iter::FromIterator;
    use euclid::Point2D;
    use super::{FovValue, HexFov};
    use hex::HexGeom;

    #[derive(PartialEq, Eq, Clone)]
    struct Cell1 {
        range: i32,
    }

    impl FovValue for Cell1 {
        fn advance(&self, offset: Point2D<i32>) -> Option<Self> {
            if offset.hex_dist() < self.range { Some(self.clone()) } else { None }
        }
    }

    #[derive(PartialEq, Eq, Clone)]
    struct Cell2 {
        range: i32,
    }

    impl FovValue for Cell2 {
        fn advance(&self, offset: Point2D<i32>) -> Option<Self> {
            if offset.hex_dist() < self.range { Some(self.clone()) } else { None }
        }

        fn is_fake_isometric_wall(&self, offset: Point2D<i32>) -> bool {
            let _ = offset;
            true
        }
    }

    #[test]
    fn trivial_fov() {
        // Just draw a small circle.
        let field: HashMap<Point2D<i32>, Cell1> = HashMap::from_iter(HexFov::new(Cell1 {
            range: 2,
        }));
        assert!(field.contains_key(&Point2D::new(1, 0)));
        assert!(!field.contains_key(&Point2D::new(1, -1)));

        // Now test out the fake-isometric corners.
        let field: HashMap<Point2D<i32>, Cell2> = HashMap::from_iter(HexFov::new(Cell2 {
            range: 2,
        }));
        assert!(field.contains_key(&Point2D::new(1, 0)));
        assert!(field.contains_key(&Point2D::new(1, -1)));
    }
}
<commit_msg>Factor hexfov corner hack into own method<commit_after>use num::Integer;
use euclid::Point2D;
use hex::Dir6;

pub trait FovValue: Eq + Clone {
    /// Construct a new FovValue based on previous one and a new point in the fov.
    fn advance(&self, offset: Point2D<i32>) -> Option<Self>;

    /// Optional method for showing acute corners of fake isometric rooms in fov.
    fn is_fake_isometric_wall(&self, offset: Point2D<i32>) -> bool {
        let _ = offset;
        false
    }
}

/// Field of view iterator for a hexagonal map.
///
/// Takes a function that maps cells into user-specified values and indicates when traversal should
/// stop.
pub struct HexFov<T> {
    stack: Vec<Sector<T>>,
    /// Extra values generated by special cases.
    side_channel: Vec<(Point2D<i32>, T)>,
}

impl<T: FovValue> HexFov<T> {
    pub fn new(init: T) -> HexFov<T> {
        // We could run f for (0, 0) here, but the traditional way for the FOV to work is to only
        // consider your surroundings, not the origin site itself.
        HexFov {
            stack: vec![Sector {
                            begin: PolarPoint::new(0.0, 1),
                            pt: PolarPoint::new(0.0, 1),
                            end: PolarPoint::new(6.0, 1),
                            prev_value: init.clone(),
                            group_value: None,
                        }],
            // The FOV algorithm will not generate the origin point, so we use
            // the side channel to explicitly add it in the beginning.
            side_channel: vec![(Point2D::new(0, 0), init)],
        }
    }

    /// Add visible horizontal corners to fake-isometric rooms.
    fn make_corners_visible(&mut self, current: &Sector<T>) {
        // We're moving along a vertical line on the hex circle, so there are side
        // points to check.
        if let Some(side_pos) = current.pt.side_point() {
            let next = current.pt.next();
            // If the next cell is within the current span and the current cell is
            // wallform,
            if next.is_below(current.end) &&
               current.prev_value.is_fake_isometric_wall(current.pt.to_v2()) {
                // and if the next cell is visible,
                if let Some(next_value) = current.prev_value.advance(next.to_v2()) {
                    // and if the current and the next cell are in the same value group,
                    // both the next cell and the third corner point cell are
                    // wallforms, and the side point would not be otherwise
                    // visible:
                    if next_value == current.prev_value &&
                       next_value.is_fake_isometric_wall(next.to_v2()) &&
                       current.prev_value.advance(side_pos).is_none() &&
                       current.prev_value.is_fake_isometric_wall(side_pos) {
                        // Add the side point to the side channel.
                        self.side_channel.push((side_pos, current.prev_value.clone()));
                    }
                }
            }
        }
    }
}

impl<T: FovValue> Iterator for HexFov<T> {
    type Item = (Point2D<i32>, T);
    fn next(&mut self) -> Option<(Point2D<i32>, T)> {
        // Empty the side channel before proceeding to the algorithm proper.
        if let Some(ret) = self.side_channel.pop() {
            return Some(ret);
        }

        // Start processing the next sector in the stack.
        if let Some(mut current) = self.stack.pop() {
            // Generate the user value from the previous value in the stack item and the working
            // position in the current sector.
            let next_value = current.prev_value.advance(current.pt.to_v2());

            // Proceed along the sector if we're not at its end yet.
            if current.pt.is_below(current.end) {
                let pos = current.pt.to_v2();

                match current.group_value {
                    None => {
                        // Beginning of group, value isn't set.
                        current.group_value = next_value.clone();
                    }
                    ref old_value if old_value != &next_value => {
                        // Value changed, branch out.

                        // Add the rest of this sector with the new value.
                        self.stack.push(Sector {
                            begin: current.pt,
                            pt: current.pt,
                            end: current.end,
                            prev_value: current.prev_value,
                            group_value: next_value.clone(),
                        });

                        // Branch further on the arc processed so far.
                        if let &Some(ref old_value) = old_value {
                            self.stack.push(Sector {
                                begin: current.begin.further(),
                                pt: current.begin.further(),
                                end: current.pt.further(),
                                prev_value: old_value.clone(),
                                group_value: None,
                            });
                        }
                        return self.next();
                    }
                    // Otherwise don't modify the stack.
                    _ => {}
                }

                // Current value and group value are assumed to be the same from here on.
                assert!(current.group_value == next_value);

                self.make_corners_visible(&current);

                // Proceed along the current sector.
                current.pt = current.pt.next();
                self.stack.push(current);
                if let Some(next_value) = next_value {
                    return Some((pos, next_value));
                } else {
                    return self.next();
                }
            } else {
                // Hit the end of the sector, branch ahead.
                if let Some(value) = current.group_value {
                    self.stack.push(Sector {
                        begin: current.begin.further(),
                        pt: current.begin.further(),
                        end: current.end.further(),
                        prev_value: value,
                        group_value: None,
                    });
                }

                self.next()
            }
        } else {
            None
        }
    }
}

struct Sector<T> {
    /// Start point of current sector.
    begin: PolarPoint,
    /// Point currently being processed.
    pt: PolarPoint,
    /// End point of current sector.
    end: PolarPoint,
    /// The user value from previous iteration.
    prev_value: T,
    /// The user value for this group.
    group_value: Option<T>,
}

/// Points on a hex circle expressed in polar coordinates.
#[derive(Copy, Clone, PartialEq)]
struct PolarPoint {
    pos: f32,
    radius: u32,
}

impl PolarPoint {
    pub fn new(pos: f32, radius: u32) -> PolarPoint {
        PolarPoint {
            pos: pos,
            radius: radius,
        }
    }
    /// Index of the discrete hex cell along the circle that corresponds to this point.
    fn winding_index(self) -> i32 { (self.pos + 0.5).floor() as i32 }

    pub fn is_below(self, other: PolarPoint) -> bool { self.winding_index() < other.end_index() }

    fn end_index(self) -> i32 { (self.pos + 0.5).ceil() as i32 }

    pub fn to_v2(self) -> Point2D<i32> {
        if self.radius == 0 {
            return Point2D::new(0, 0);
        }
        let index = self.winding_index();
        let sector = index.mod_floor(&(self.radius as i32 * 6)) / self.radius as i32;
        let offset = index.mod_floor(&(self.radius as i32));

        let rod = Dir6::from_int(sector).to_v2();
        let tangent = Dir6::from_int((sector + 2) % 6).to_v2();

        rod * (self.radius as i32) + tangent * offset
    }

    /// If this point and the next point are adjacent vertically (along the xy
    /// axis), return the point outside of the circle between the two points.
    ///
    /// This is a helper function for the FOV special case where acute corners
    /// of fake isometric rooms are marked visible even though strict hex FOV
    /// logic would keep them unseen.
    pub fn side_point(self) -> Option<Point2D<i32>> {
        let next = self.next();
        let a = self.to_v2();
        let b = next.to_v2();

        if b.x == a.x + 1 && b.y == a.y + 1 {
            // Going down the right rim.
            Some(Point2D::new(a.x + 1, a.y))
        } else if b.x == a.x - 1 && b.y == a.y - 1 {
            // Going up the left rim.
            Some(Point2D::new(a.x - 1, a.y))
        } else {
            None
        }
    }

    /// The point corresponding to this one on the hex circle with radius +1.
    pub fn further(self) -> PolarPoint {
        PolarPoint::new(self.pos * (self.radius + 1) as f32 / self.radius as f32,
                        self.radius + 1)
    }

    /// The point next to this one along the hex circle.
    pub fn next(self) -> PolarPoint { PolarPoint::new((self.pos + 0.5).floor() + 0.5, self.radius) }
}

#[cfg(test)]
mod test {
    use std::collections::HashMap;
    use std::iter::FromIterator;
    use euclid::Point2D;
    use super::{FovValue, HexFov};
    use hex::HexGeom;

    #[derive(PartialEq, Eq, Clone)]
    struct Cell1 {
        range: i32,
    }

    impl FovValue for Cell1 {
        fn advance(&self, offset: Point2D<i32>) -> Option<Self> {
            if offset.hex_dist() < self.range { Some(self.clone()) } else { None }
        }
    }

    #[derive(PartialEq, Eq, Clone)]
    struct Cell2 {
        range: i32,
    }

    impl FovValue for Cell2 {
        fn advance(&self, offset: Point2D<i32>) -> Option<Self> {
            if offset.hex_dist() < self.range { Some(self.clone()) } else { None }
        }

        fn is_fake_isometric_wall(&self, offset: Point2D<i32>) -> bool {
            let _ = offset;
            true
        }
    }

    #[test]
    fn trivial_fov() {
        // Just draw a small circle.
        let field: HashMap<Point2D<i32>, Cell1> = HashMap::from_iter(HexFov::new(Cell1 {
            range: 2,
        }));
        assert!(field.contains_key(&Point2D::new(1, 0)));
        assert!(!field.contains_key(&Point2D::new(1, -1)));

        // Now test out the fake-isometric corners.
        let field: HashMap<Point2D<i32>, Cell2> = HashMap::from_iter(HexFov::new(Cell2 {
            range: 2,
        }));
        assert!(field.contains_key(&Point2D::new(1, 0)));
        assert!(field.contains_key(&Point2D::new(1, -1)));
    }
}
<|endoftext|>