<commit_before>// Copyright 2014-2015 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 io::{self, Error, ErrorKind};
use libc::{self, c_int, gid_t, pid_t, uid_t};
use ptr;

use sys::cvt;
use sys::process::process_common::*;

////////////////////////////////////////////////////////////////////////////////
// Command
////////////////////////////////////////////////////////////////////////////////

impl Command {
    pub fn spawn(&mut self, default: Stdio, needs_stdin: bool)
                 -> io::Result<(Process, StdioPipes)> {
        use sys;

        const CLOEXEC_MSG_FOOTER: &'static [u8] = b"NOEX";

        let envp = self.capture_env();

        if self.saw_nul() {
            return Err(io::Error::new(ErrorKind::InvalidInput,
                                      "nul byte found in provided data"));
        }

        let (ours, theirs) = self.setup_io(default, needs_stdin)?;
        let (input, output) = sys::pipe::anon_pipe()?;

        let pid = unsafe {
            match cvt(libc::fork())? {
                0 => {
                    drop(input);
                    let err = self.do_exec(theirs, envp.as_ref());
                    let errno = err.raw_os_error().unwrap_or(libc::EINVAL) as u32;
                    let bytes = [
                        (errno >> 24) as u8,
                        (errno >> 16) as u8,
                        (errno >>  8) as u8,
                        (errno >>  0) as u8,
                        CLOEXEC_MSG_FOOTER[0], CLOEXEC_MSG_FOOTER[1],
                        CLOEXEC_MSG_FOOTER[2], CLOEXEC_MSG_FOOTER[3]
                    ];
                    // pipe I/O up to PIPE_BUF bytes should be atomic, and then
                    // we want to be sure we *don't* run at_exit destructors as
                    // we're being torn down regardless
                    assert!(output.write(&bytes).is_ok());
                    libc::_exit(1)
                }
                n => n,
            }
        };

        let mut p = Process { pid: pid, status: None };
        drop(output);
        let mut bytes = [0; 8];

        // loop to handle EINTR
        loop {
            match input.read(&mut bytes) {
                Ok(0) => return Ok((p, ours)),
                Ok(8) => {
                    assert!(combine(CLOEXEC_MSG_FOOTER) == combine(&bytes[4.. 8]),
                            "Validation on the CLOEXEC pipe failed: {:?}", bytes);
                    let errno = combine(&bytes[0.. 4]);
                    assert!(p.wait().is_ok(),
                            "wait() should either return Ok or panic");
                    return Err(Error::from_raw_os_error(errno))
                }
                Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
                Err(e) => {
                    assert!(p.wait().is_ok(),
                            "wait() should either return Ok or panic");
                    panic!("the CLOEXEC pipe failed: {:?}", e)
                },
                Ok(..) => { // pipe I/O up to PIPE_BUF bytes should be atomic
                    assert!(p.wait().is_ok(),
                            "wait() should either return Ok or panic");
                    panic!("short read on the CLOEXEC pipe")
                }
            }
        }

        fn combine(arr: &[u8]) -> i32 {
            let a = arr[0] as u32;
            let b = arr[1] as u32;
            let c = arr[2] as u32;
            let d = arr[3] as u32;

            ((a << 24) | (b << 16) | (c << 8) | (d << 0)) as i32
        }
    }

    pub fn exec(&mut self, default: Stdio) -> io::Error {
        let envp = self.capture_env();

        if self.saw_nul() {
            return io::Error::new(ErrorKind::InvalidInput,
                                  "nul byte found in provided data")
        }

        match self.setup_io(default, true) {
            Ok((_, theirs)) => unsafe { self.do_exec(theirs, envp.as_ref()) },
            Err(e) => e,
        }
    }

    // And at this point we've reached a special time in the life of the
    // child. The child must now be considered hamstrung and unable to
    // do anything other than syscalls really. Consider the following
    // scenario:
    //
    //      1. Thread A of process 1 grabs the malloc() mutex
    //      2. Thread B of process 1 forks(), creating thread C
    //      3. Thread C of process 2 then attempts to malloc()
    //      4. The memory of process 2 is the same as the memory of
    //         process 1, so the mutex is locked.
    //
    // This situation looks a lot like deadlock, right? It turns out
    // that this is what pthread_atfork() takes care of, which is
    // presumably implemented across platforms. The first thing that
    // threads to *before* forking is to do things like grab the malloc
    // mutex, and then after the fork they unlock it.
    //
    // Despite this information, libnative's spawn has been witnessed to
    // deadlock on both macOS and FreeBSD. I'm not entirely sure why, but
    // all collected backtraces point at malloc/free traffic in the
    // child spawned process.
    //
    // For this reason, the block of code below should contain 0
    // invocations of either malloc of free (or their related friends).
    //
    // As an example of not having malloc/free traffic, we don't close
    // this file descriptor by dropping the FileDesc (which contains an
    // allocation). Instead we just close it manually. This will never
    // have the drop glue anyway because this code never returns (the
    // child will either exec() or invoke libc::exit)
    unsafe fn do_exec(
        &mut self,
        stdio: ChildPipes,
        maybe_envp: Option<&CStringArray>
    ) -> io::Error {
        use sys::{self, cvt_r};

        macro_rules! t {
            ($e:expr) => (match $e {
                Ok(e) => e,
                Err(e) => return e,
            })
        }

        if let Some(fd) = stdio.stdin.fd() {
            t!(cvt_r(|| libc::dup2(fd, libc::STDIN_FILENO)));
        }
        if let Some(fd) = stdio.stdout.fd() {
            t!(cvt_r(|| libc::dup2(fd, libc::STDOUT_FILENO)));
        }
        if let Some(fd) = stdio.stderr.fd() {
            t!(cvt_r(|| libc::dup2(fd, libc::STDERR_FILENO)));
        }

        if cfg!(not(any(target_os = "l4re"))) {
            if let Some(u) = self.get_gid() {
                t!(cvt(libc::setgid(u as gid_t)));
            }
            if let Some(u) = self.get_uid() {
                // When dropping privileges from root, the `setgroups` call
                // will remove any extraneous groups. If we don't call this,
                // then even though our uid has dropped, we may still have
                // groups that enable us to do super-user things. This will
                // fail if we aren't root, so don't bother checking the
                // return value, this is just done as an optimistic
                // privilege dropping function.
                let _ = libc::setgroups(0, ptr::null());

                t!(cvt(libc::setuid(u as uid_t)));
            }
        }
        if let Some(ref cwd) = *self.get_cwd() {
            t!(cvt(libc::chdir(cwd.as_ptr())));
        }
        if let Some(envp) = maybe_envp {
            *sys::os::environ() = envp.as_ptr();
        }

        // emscripten has no signal support.
        #[cfg(not(any(target_os = "emscripten")))]
        {
            use mem;
            // Reset signal handling so the child process starts in a
            // standardized state. libstd ignores SIGPIPE, and signal-handling
            // libraries often set a mask. Child processes inherit ignored
            // signals and the signal mask from their parent, but most
            // UNIX programs do not reset these things on their own, so we
            // need to clean things up now to avoid confusing the program
            // we're about to run.
            let mut set: libc::sigset_t = mem::uninitialized();
            if cfg!(target_os = "android") {
                // Implementing sigemptyset allow us to support older Android
                // versions. See the comment about Android and sig* functions in
                // process_common.rs
                libc::memset(&mut set as *mut _ as *mut _,
                             0,
                             mem::size_of::<libc::sigset_t>());
            } else {
                t!(cvt(libc::sigemptyset(&mut set)));
            }
            t!(cvt(libc::pthread_sigmask(libc::SIG_SETMASK, &set,
                                         ptr::null_mut())));
            let ret = sys::signal(libc::SIGPIPE, libc::SIG_DFL);
            if ret == libc::SIG_ERR {
                return io::Error::last_os_error()
            }
        }

        for callback in self.get_closures().iter_mut() {
            t!(callback());
        }

        libc::execvp(self.get_argv()[0], self.get_argv().as_ptr());
        io::Error::last_os_error()
    }
}

////////////////////////////////////////////////////////////////////////////////
// Processes
////////////////////////////////////////////////////////////////////////////////

/// The unique id of the process (this should never be negative).
pub struct Process {
    pid: pid_t,
    status: Option<ExitStatus>,
}

impl Process {
    pub fn id(&self) -> u32 {
        self.pid as u32
    }

    pub fn kill(&mut self) -> io::Result<()> {
        // If we've already waited on this process then the pid can be recycled
        // and used for another process, and we probably shouldn't be killing
        // random processes, so just return an error.
        if self.status.is_some() {
            Err(Error::new(ErrorKind::InvalidInput,
                           "invalid argument: can't kill an exited process"))
        } else {
            cvt(unsafe { libc::kill(self.pid, libc::SIGKILL) }).map(|_| ())
        }
    }

    pub fn wait(&mut self) -> io::Result<ExitStatus> {
        use sys::cvt_r;
        if let Some(status) = self.status {
            return Ok(status)
        }
        let mut status = 0 as c_int;
        cvt_r(|| unsafe { libc::waitpid(self.pid, &mut status, 0) })?;
        self.status = Some(ExitStatus::new(status));
        Ok(ExitStatus::new(status))
    }

    pub fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
        if let Some(status) = self.status {
            return Ok(Some(status))
        }
        let mut status = 0 as c_int;
        let pid = cvt(unsafe {
            libc::waitpid(self.pid, &mut status, libc::WNOHANG)
        })?;
        if pid == 0 {
            Ok(None)
        } else {
            self.status = Some(ExitStatus::new(status));
            Ok(Some(ExitStatus::new(status)))
        }
    }
}
<commit_msg>Support posix_spawn() when possible.<commit_after>// Copyright 2014-2015 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 io::{self, Error, ErrorKind};
use libc::{self, c_int, gid_t, pid_t, uid_t};
use ptr;

use sys::cvt;
use sys::process::process_common::*;

////////////////////////////////////////////////////////////////////////////////
// Command
////////////////////////////////////////////////////////////////////////////////

impl Command {
    pub fn spawn(&mut self, default: Stdio, needs_stdin: bool)
                 -> io::Result<(Process, StdioPipes)> {
        use sys;

        const CLOEXEC_MSG_FOOTER: &'static [u8] = b"NOEX";

        let envp = self.capture_env();

        if self.saw_nul() {
            return Err(io::Error::new(ErrorKind::InvalidInput,
                                      "nul byte found in provided data"));
        }

        let (ours, theirs) = self.setup_io(default, needs_stdin)?;

        if let Some(ret) = self.posix_spawn(&theirs, envp.as_ref())? {
            return Ok((ret, ours))
        }

        let (input, output) = sys::pipe::anon_pipe()?;

        let pid = unsafe {
            match cvt(libc::fork())? {
                0 => {
                    drop(input);
                    let err = self.do_exec(theirs, envp.as_ref());
                    let errno = err.raw_os_error().unwrap_or(libc::EINVAL) as u32;
                    let bytes = [
                        (errno >> 24) as u8,
                        (errno >> 16) as u8,
                        (errno >>  8) as u8,
                        (errno >>  0) as u8,
                        CLOEXEC_MSG_FOOTER[0], CLOEXEC_MSG_FOOTER[1],
                        CLOEXEC_MSG_FOOTER[2], CLOEXEC_MSG_FOOTER[3]
                    ];
                    // pipe I/O up to PIPE_BUF bytes should be atomic, and then
                    // we want to be sure we *don't* run at_exit destructors as
                    // we're being torn down regardless
                    assert!(output.write(&bytes).is_ok());
                    libc::_exit(1)
                }
                n => n,
            }
        };

        let mut p = Process { pid: pid, status: None };
        drop(output);
        let mut bytes = [0; 8];

        // loop to handle EINTR
        loop {
            match input.read(&mut bytes) {
                Ok(0) => return Ok((p, ours)),
                Ok(8) => {
                    assert!(combine(CLOEXEC_MSG_FOOTER) == combine(&bytes[4.. 8]),
                            "Validation on the CLOEXEC pipe failed: {:?}", bytes);
                    let errno = combine(&bytes[0.. 4]);
                    assert!(p.wait().is_ok(),
                            "wait() should either return Ok or panic");
                    return Err(Error::from_raw_os_error(errno))
                }
                Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
                Err(e) => {
                    assert!(p.wait().is_ok(),
                            "wait() should either return Ok or panic");
                    panic!("the CLOEXEC pipe failed: {:?}", e)
                },
                Ok(..) => { // pipe I/O up to PIPE_BUF bytes should be atomic
                    assert!(p.wait().is_ok(),
                            "wait() should either return Ok or panic");
                    panic!("short read on the CLOEXEC pipe")
                }
            }
        }

        fn combine(arr: &[u8]) -> i32 {
            let a = arr[0] as u32;
            let b = arr[1] as u32;
            let c = arr[2] as u32;
            let d = arr[3] as u32;

            ((a << 24) | (b << 16) | (c << 8) | (d << 0)) as i32
        }
    }

    pub fn exec(&mut self, default: Stdio) -> io::Error {
        let envp = self.capture_env();

        if self.saw_nul() {
            return io::Error::new(ErrorKind::InvalidInput,
                                  "nul byte found in provided data")
        }

        match self.setup_io(default, true) {
            Ok((_, theirs)) => unsafe { self.do_exec(theirs, envp.as_ref()) },
            Err(e) => e,
        }
    }

    // And at this point we've reached a special time in the life of the
    // child. The child must now be considered hamstrung and unable to
    // do anything other than syscalls really. Consider the following
    // scenario:
    //
    //      1. Thread A of process 1 grabs the malloc() mutex
    //      2. Thread B of process 1 forks(), creating thread C
    //      3. Thread C of process 2 then attempts to malloc()
    //      4. The memory of process 2 is the same as the memory of
    //         process 1, so the mutex is locked.
    //
    // This situation looks a lot like deadlock, right? It turns out
    // that this is what pthread_atfork() takes care of, which is
    // presumably implemented across platforms. The first thing that
    // threads to *before* forking is to do things like grab the malloc
    // mutex, and then after the fork they unlock it.
    //
    // Despite this information, libnative's spawn has been witnessed to
    // deadlock on both macOS and FreeBSD. I'm not entirely sure why, but
    // all collected backtraces point at malloc/free traffic in the
    // child spawned process.
    //
    // For this reason, the block of code below should contain 0
    // invocations of either malloc of free (or their related friends).
    //
    // As an example of not having malloc/free traffic, we don't close
    // this file descriptor by dropping the FileDesc (which contains an
    // allocation). Instead we just close it manually. This will never
    // have the drop glue anyway because this code never returns (the
    // child will either exec() or invoke libc::exit)
    unsafe fn do_exec(
        &mut self,
        stdio: ChildPipes,
        maybe_envp: Option<&CStringArray>
    ) -> io::Error {
        use sys::{self, cvt_r};

        macro_rules! t {
            ($e:expr) => (match $e {
                Ok(e) => e,
                Err(e) => return e,
            })
        }

        if let Some(fd) = stdio.stdin.fd() {
            t!(cvt_r(|| libc::dup2(fd, libc::STDIN_FILENO)));
        }
        if let Some(fd) = stdio.stdout.fd() {
            t!(cvt_r(|| libc::dup2(fd, libc::STDOUT_FILENO)));
        }
        if let Some(fd) = stdio.stderr.fd() {
            t!(cvt_r(|| libc::dup2(fd, libc::STDERR_FILENO)));
        }

        if cfg!(not(any(target_os = "l4re"))) {
            if let Some(u) = self.get_gid() {
                t!(cvt(libc::setgid(u as gid_t)));
            }
            if let Some(u) = self.get_uid() {
                // When dropping privileges from root, the `setgroups` call
                // will remove any extraneous groups. If we don't call this,
                // then even though our uid has dropped, we may still have
                // groups that enable us to do super-user things. This will
                // fail if we aren't root, so don't bother checking the
                // return value, this is just done as an optimistic
                // privilege dropping function.
                let _ = libc::setgroups(0, ptr::null());

                t!(cvt(libc::setuid(u as uid_t)));
            }
        }
        if let Some(ref cwd) = *self.get_cwd() {
            t!(cvt(libc::chdir(cwd.as_ptr())));
        }
        if let Some(envp) = maybe_envp {
            *sys::os::environ() = envp.as_ptr();
        }

        // emscripten has no signal support.
        #[cfg(not(any(target_os = "emscripten")))]
        {
            use mem;
            // Reset signal handling so the child process starts in a
            // standardized state. libstd ignores SIGPIPE, and signal-handling
            // libraries often set a mask. Child processes inherit ignored
            // signals and the signal mask from their parent, but most
            // UNIX programs do not reset these things on their own, so we
            // need to clean things up now to avoid confusing the program
            // we're about to run.
            let mut set: libc::sigset_t = mem::uninitialized();
            if cfg!(target_os = "android") {
                // Implementing sigemptyset allow us to support older Android
                // versions. See the comment about Android and sig* functions in
                // process_common.rs
                libc::memset(&mut set as *mut _ as *mut _,
                             0,
                             mem::size_of::<libc::sigset_t>());
            } else {
                t!(cvt(libc::sigemptyset(&mut set)));
            }
            t!(cvt(libc::pthread_sigmask(libc::SIG_SETMASK, &set,
                                         ptr::null_mut())));
            let ret = sys::signal(libc::SIGPIPE, libc::SIG_DFL);
            if ret == libc::SIG_ERR {
                return io::Error::last_os_error()
            }
        }

        for callback in self.get_closures().iter_mut() {
            t!(callback());
        }

        libc::execvp(self.get_argv()[0], self.get_argv().as_ptr());
        io::Error::last_os_error()
    }

    #[cfg(not(any(target_os = "linux", target_os = "macos")))]
    fn posix_spawn(&mut self, stdio: &ChildPipes, envp: Option<&CStringArray>)
        -> io::Result<Option<Process>>
    {
        Ok(None)
    }

    #[cfg(any(target_os = "linux", target_os = "macos"))]
    fn posix_spawn(&mut self, stdio: &ChildPipes, envp: Option<&CStringArray>)
        -> io::Result<Option<Process>>
    {
        use mem;
        use sys;

        if self.get_cwd().is_some() ||
            self.get_gid().is_some() ||
            self.get_uid().is_some() ||
            self.get_closures().len() != 0 {
            return Ok(None)
        }

        let mut p = Process { pid: 0, status: None };

        struct PosixSpawnFileActions(libc::posix_spawn_file_actions_t);

        impl Drop for PosixSpawnFileActions {
            fn drop(&mut self) {
                unsafe {
                    libc::posix_spawn_file_actions_destroy(&mut self.0);
                }
            }
        }

        struct PosixSpawnattr(libc::posix_spawnattr_t);

        impl Drop for PosixSpawnattr {
            fn drop(&mut self) {
                unsafe {
                    libc::posix_spawnattr_destroy(&mut self.0);
                }
            }
        }

        unsafe {
            let mut file_actions = PosixSpawnFileActions(mem::zeroed());
            let mut attrs = PosixSpawnattr(mem::zeroed());

            libc::posix_spawnattr_init(&mut attrs.0);
            libc::posix_spawn_file_actions_init(&mut file_actions.0);

            if let Some(fd) = stdio.stdin.fd() {
                cvt(libc::posix_spawn_file_actions_adddup2(&mut file_actions.0,
                                                           fd,
                                                           libc::STDIN_FILENO))?;
            }
            if let Some(fd) = stdio.stdout.fd() {
                cvt(libc::posix_spawn_file_actions_adddup2(&mut file_actions.0,
                                                           fd,
                                                           libc::STDOUT_FILENO))?;
            }
            if let Some(fd) = stdio.stderr.fd() {
                cvt(libc::posix_spawn_file_actions_adddup2(&mut file_actions.0,
                                                           fd,
                                                           libc::STDERR_FILENO))?;
            }

            let mut set: libc::sigset_t = mem::zeroed();
            cvt(libc::sigemptyset(&mut set))?;
            cvt(libc::posix_spawnattr_setsigmask(&mut attrs.0,
                                                 &set))?;
            cvt(libc::sigaddset(&mut set, libc::SIGPIPE))?;
            cvt(libc::posix_spawnattr_setsigdefault(&mut attrs.0,
                                                    &set))?;

            let flags = libc::POSIX_SPAWN_SETSIGDEF |
                libc::POSIX_SPAWN_SETSIGMASK;
            cvt(libc::posix_spawnattr_setflags(&mut attrs.0, flags as _))?;

            let envp = envp.map(|c| c.as_ptr())
                .unwrap_or(sys::os::environ() as *const _);
            let ret = libc::posix_spawnp(
                &mut p.pid,
                self.get_argv()[0],
                &file_actions.0,
                &attrs.0,
                self.get_argv().as_ptr() as *const _,
                envp as *const _,
            );
            if ret == 0 {
                Ok(Some(p))
            } else {
                Err(io::Error::last_os_error())
            }
        }
    }
}

////////////////////////////////////////////////////////////////////////////////
// Processes
////////////////////////////////////////////////////////////////////////////////

/// The unique id of the process (this should never be negative).
pub struct Process {
    pid: pid_t,
    status: Option<ExitStatus>,
}

impl Process {
    pub fn id(&self) -> u32 {
        self.pid as u32
    }

    pub fn kill(&mut self) -> io::Result<()> {
        // If we've already waited on this process then the pid can be recycled
        // and used for another process, and we probably shouldn't be killing
        // random processes, so just return an error.
        if self.status.is_some() {
            Err(Error::new(ErrorKind::InvalidInput,
                           "invalid argument: can't kill an exited process"))
        } else {
            cvt(unsafe { libc::kill(self.pid, libc::SIGKILL) }).map(|_| ())
        }
    }

    pub fn wait(&mut self) -> io::Result<ExitStatus> {
        use sys::cvt_r;
        if let Some(status) = self.status {
            return Ok(status)
        }
        let mut status = 0 as c_int;
        cvt_r(|| unsafe { libc::waitpid(self.pid, &mut status, 0) })?;
        self.status = Some(ExitStatus::new(status));
        Ok(ExitStatus::new(status))
    }

    pub fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
        if let Some(status) = self.status {
            return Ok(Some(status))
        }
        let mut status = 0 as c_int;
        let pid = cvt(unsafe {
            libc::waitpid(self.pid, &mut status, libc::WNOHANG)
        })?;
        if pid == 0 {
            Ok(None)
        } else {
            self.status = Some(ExitStatus::new(status));
            Ok(Some(ExitStatus::new(status)))
        }
    }
}
<|endoftext|>


<commit_before>#![allow(deprecated)]
use std::path::PathBuf;
use std::{io, fs, slice};
use curl::easy::Easy;
use memmap::{Mmap, Protection};
use std::io::{Read, Write};
use utils::data::{DatasetIntfRef, DatasetIntf};
use std::rc::Rc;
use tensor::{Tensor, TensorKind, NumLimits};
use torch;
use std::marker::PhantomData;

type Sample<T> = (Tensor<T>, i64);
type CollatedSample<T> = (Tensor<T>, Tensor<i64>);

static URLS: [&str; 4] = ["http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz"];
static RAW_FOLDER: &str = "raw";
static PROCESSED_FOLDER: &str = "processed";
static TRAINING_FILE: &str = "training.pt";
static TEST_FILE: &str = "test.pt";
static NCHANNELS: isize = 1;

fn create_dir_f(arg: PathBuf) -> io::Result<()> {
    let result = fs::create_dir(arg);
    if let Some(err) = result.err() {
        match err.kind() {
            io::ErrorKind::AlreadyExists => Ok(()),
            _ => Err(err),
        }
    } else {
        Ok(())
    }
}

fn check_exists(arg: &PathBuf) -> io::Result<bool> {
    let train = fs::metadata(arg.join(TRAINING_FILE))?;
    let test = fs::metadata(arg.join(TEST_FILE))?;

    return Ok(train.is_file() && test.is_file());
}

fn download(root: &String) -> io::Result<()> {
    let raw_path = PathBuf::from(root).join(RAW_FOLDER);
    let processed_path = PathBuf::from(root).join(PROCESSED_FOLDER);
    create_dir_f(raw_path.clone())?;
    create_dir_f(processed_path.clone())?;

    for url in URLS.iter() {
        println!("downloading {}", url);
        let paths: Vec<_> = url.split("/").collect();
        let mut fname: Vec<_> = paths[5].split("/").collect();
        fname = fname[0].split(".").collect();
        let file_path = raw_path.join(fname[0]);
        let mut file = fs::OpenOptions::new()
            .write(true)
            .create(true)
            .open(file_path)?;
        let mut data = Vec::new();
        let mut handle = Easy::new();
        handle.url(url).unwrap();
        {
            let mut transfer = handle.transfer();
            transfer
                .write_function(|new_data| {
                                    data.extend_from_slice(new_data);
                                    Ok(new_data.len())
                                })
                .unwrap();
            transfer.perform().unwrap();
        }
        let mut output = Vec::with_capacity(data.len() + 1100000);
        let mut gz = ::flate2::read::GzDecoder::new(data.as_slice())?;
        gz.read_to_end(&mut output)?;

        file.write_all(output.as_slice())?;
    }
    println!("Proceeding");

    let training_set =
        vec![read_image_file(raw_path.join("train-images-idx3-ubyte"))?,
                            read_label_file(raw_path.join("train-labels-idx1-ubyte"))?];
    let test_set = vec![read_image_file(raw_path.join("t10k-images-idx3-ubyte"))?,
                        read_label_file(raw_path.join("t10k-labels-idx1-ubyte"))?];

    torch::save(processed_path.join(TRAINING_FILE), &training_set)?;
    torch::save(processed_path.join(TEST_FILE), &test_set)?;

    Ok(())
}

fn read_image_file(path: PathBuf) -> io::Result<TensorKind> {
    let fp = Mmap::open_path(path, Protection::Read)?;
    let (length, nrows, ncols);
    {
        let data = unsafe { slice::from_raw_parts(fp.ptr() as *const i32, fp.len()) };
        assert_eq!(i32::from_be(data[0]), 2051);
        length = i32::from_be(data[1]) as usize;
        nrows = i32::from_be(data[2]) as usize;
        ncols = i32::from_be(data[3]) as usize;
    }
    let data = unsafe { fp.as_slice() };
    let mut images: Vec<Vec<u8>> = Vec::with_capacity(length);
    let mut idx = 16;
    // XXX transpose here to be H x W ?
    for _ in 0..length {
        let mut img = Vec::with_capacity(nrows * ncols);
        for _ in 0..nrows * ncols {
            img.push(data[idx]);
            idx += 1;
        }
        images.push(img)
    }
    Ok(torch::byte_tensor(images).view(&[-1, NCHANNELS, 28, 28]).into())
}

fn read_label_file(path: PathBuf) -> io::Result<TensorKind> {
    let fp = Mmap::open_path(path, Protection::Read)?;
    let length;
    {
        let data = unsafe { slice::from_raw_parts(fp.ptr() as *const i32, fp.len()) };
        assert_eq!(i32::from_be(data[0]), 2049);
        length = i32::from_be(data[1]) as usize;
    }
    let data = unsafe { fp.as_slice() };
    let mut labels: Vec<u8> = Vec::with_capacity(length);
    for i in 0..length {
        labels.push(data[8 + i])
    }
    Ok(torch::byte_tensor(labels).into())
}

pub struct MNIST<T: NumLimits> {
    pub root: String,
    pub train: bool,
    pub data: Tensor<u8>,
    pub labels: Tensor<u8>,
    pub transform: Option<Xfrm>,
    phantom: PhantomData<T>,
}

#[builder(pattern="owned")]
#[derive(Builder)]
pub struct MNISTArgs {
    root: String,
    #[builder(default="true")]
    train: bool,
    #[builder(default="false")]
    download: bool,
}

type Xfrm = Box<fn(&TensorKind) -> TensorKind>;
impl MNISTArgsBuilder {
    pub fn done<T: NumLimits + 'static>(self,
                                        xfrm: Option<Xfrm>)
                                        -> DatasetIntfRef<CollatedSample<T>> {
        let args = self.build().unwrap();
        Rc::new(MNIST::new(args, xfrm))
    }
}

impl<T: NumLimits> MNIST<T> {
    pub fn build(root: &str) -> MNISTArgsBuilder {
        MNISTArgsBuilder::default().root(root.into())
    }
    pub fn new(args: MNISTArgs, xfrm: Option<Xfrm>) -> Self {
        if args.download {
            download(&args.root).expect("download failed");
        }
        let processed_path = PathBuf::from(&args.root).join(PROCESSED_FOLDER);
        if !check_exists(&processed_path).expect("Dataset not found, try downloading") {
            panic!("Dataset not found, use download=true to download it");
        }
        let mut v: Vec<TensorKind> = if args.train {
            torch::load(processed_path.join(TRAINING_FILE)).expect("torch load failed")
        } else {
            torch::load(processed_path.join(TEST_FILE)).expect("torch load failed")
        };
        let (data, labels) = (v.remove(0), v.remove(0));

        MNIST {
            root: args.root.clone(),
            train: args.train,
            data: data.into(),
            labels: labels.into(),
            transform: xfrm,
            phantom: PhantomData,
        }
    }
    fn index(&self, idx: usize) -> Sample<u8> {
        let img = self.data.s([idx]);
        let img = if let Some(ref transform) = self.transform {
            transform(&img.into()).into()
        } else {
            img.clone()
        };
        (img, self.labels[idx].clone() as i64)
    }
}

impl<T: NumLimits> DatasetIntf for MNIST<T> {
    type Batch = CollatedSample<T>;
    fn len(&self) -> usize {
        if self.train { 60000 } else { 10000 }
    }
    fn collate(&self, sample: Vec<usize>) -> Self::Batch {
        println!("collate called len: {}", sample.len());
        let v: Vec<Sample<u8>> = sample.into_iter().map(|i| self.index(i)).collect();
        let labels: Vec<i64> = v.iter().map(|&(_, ref t)| *t).collect();
        let imgs: Vec<Tensor<u8>> = v.into_iter().map(|(d, _)| d).collect();
        // XXX should check for case of double
        let img_batch: Tensor<T> = torch::tensor(imgs);
        let label_batch = torch::long_tensor(labels);
        (img_batch, label_batch)
    }
}
<commit_msg>rustfmt<commit_after>#![allow(deprecated)]
use std::path::PathBuf;
use std::{io, fs, slice};
use curl::easy::Easy;
use memmap::{Mmap, Protection};
use std::io::{Read, Write};
use utils::data::{DatasetIntfRef, DatasetIntf};
use std::rc::Rc;
use tensor::{Tensor, TensorKind, NumLimits};
use torch;
use std::marker::PhantomData;

type Sample<T> = (Tensor<T>, i64);
type CollatedSample<T> = (Tensor<T>, Tensor<i64>);

static URLS: [&str; 4] = ["http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz"];
static RAW_FOLDER: &str = "raw";
static PROCESSED_FOLDER: &str = "processed";
static TRAINING_FILE: &str = "training.pt";
static TEST_FILE: &str = "test.pt";
static NCHANNELS: isize = 1;

fn create_dir_f(arg: PathBuf) -> io::Result<()> {
    let result = fs::create_dir(arg);
    if let Some(err) = result.err() {
        match err.kind() {
            io::ErrorKind::AlreadyExists => Ok(()),
            _ => Err(err),
        }
    } else {
        Ok(())
    }
}

fn check_exists(arg: &PathBuf) -> io::Result<bool> {
    let train = fs::metadata(arg.join(TRAINING_FILE))?;
    let test = fs::metadata(arg.join(TEST_FILE))?;

    return Ok(train.is_file() && test.is_file());
}

fn download(root: &String) -> io::Result<()> {
    let raw_path = PathBuf::from(root).join(RAW_FOLDER);
    let processed_path = PathBuf::from(root).join(PROCESSED_FOLDER);
    create_dir_f(raw_path.clone())?;
    create_dir_f(processed_path.clone())?;

    for url in URLS.iter() {
        println!("downloading {}", url);
        let paths: Vec<_> = url.split("/").collect();
        let mut fname: Vec<_> = paths[5].split("/").collect();
        fname = fname[0].split(".").collect();
        let file_path = raw_path.join(fname[0]);
        let mut file = fs::OpenOptions::new()
            .write(true)
            .create(true)
            .open(file_path)?;
        let mut data = Vec::new();
        let mut handle = Easy::new();
        handle.url(url).unwrap();
        {
            let mut transfer = handle.transfer();
            transfer
                .write_function(|new_data| {
                                    data.extend_from_slice(new_data);
                                    Ok(new_data.len())
                                })
                .unwrap();
            transfer.perform().unwrap();
        }
        let mut output = Vec::with_capacity(data.len() + 1100000);
        let mut gz = ::flate2::read::GzDecoder::new(data.as_slice())?;
        gz.read_to_end(&mut output)?;

        file.write_all(output.as_slice())?;
    }
    println!("Proceeding");

    let training_set =
        vec![read_image_file(raw_path.join("train-images-idx3-ubyte"))?,
                            read_label_file(raw_path.join("train-labels-idx1-ubyte"))?];
    let test_set = vec![read_image_file(raw_path.join("t10k-images-idx3-ubyte"))?,
                        read_label_file(raw_path.join("t10k-labels-idx1-ubyte"))?];

    torch::save(processed_path.join(TRAINING_FILE), &training_set)?;
    torch::save(processed_path.join(TEST_FILE), &test_set)?;

    Ok(())
}

fn read_image_file(path: PathBuf) -> io::Result<TensorKind> {
    let fp = Mmap::open_path(path, Protection::Read)?;
    let (length, nrows, ncols);
    {
        let data = unsafe { slice::from_raw_parts(fp.ptr() as *const i32, fp.len()) };
        assert_eq!(i32::from_be(data[0]), 2051);
        length = i32::from_be(data[1]) as usize;
        nrows = i32::from_be(data[2]) as usize;
        ncols = i32::from_be(data[3]) as usize;
    }
    let data = unsafe { fp.as_slice() };
    let mut images: Vec<Vec<u8>> = Vec::with_capacity(length);
    let mut idx = 16;
    // XXX transpose here to be H x W ?
    for _ in 0..length {
        let mut img = Vec::with_capacity(nrows * ncols);
        for _ in 0..nrows * ncols {
            img.push(data[idx]);
            idx += 1;
        }
        images.push(img)
    }
    Ok(torch::byte_tensor(images)
           .view(&[-1, NCHANNELS, 28, 28])
           .into())
}

fn read_label_file(path: PathBuf) -> io::Result<TensorKind> {
    let fp = Mmap::open_path(path, Protection::Read)?;
    let length;
    {
        let data = unsafe { slice::from_raw_parts(fp.ptr() as *const i32, fp.len()) };
        assert_eq!(i32::from_be(data[0]), 2049);
        length = i32::from_be(data[1]) as usize;
    }
    let data = unsafe { fp.as_slice() };
    let mut labels: Vec<u8> = Vec::with_capacity(length);
    for i in 0..length {
        labels.push(data[8 + i])
    }
    Ok(torch::byte_tensor(labels).into())
}

pub struct MNIST<T: NumLimits> {
    pub root: String,
    pub train: bool,
    pub data: Tensor<u8>,
    pub labels: Tensor<u8>,
    pub transform: Option<Xfrm>,
    phantom: PhantomData<T>,
}

#[builder(pattern="owned")]
#[derive(Builder)]
pub struct MNISTArgs {
    root: String,
    #[builder(default="true")]
    train: bool,
    #[builder(default="false")]
    download: bool,
}

type Xfrm = Box<fn(&TensorKind) -> TensorKind>;
impl MNISTArgsBuilder {
    pub fn done<T: NumLimits + 'static>(self,
                                        xfrm: Option<Xfrm>)
                                        -> DatasetIntfRef<CollatedSample<T>> {
        let args = self.build().unwrap();
        Rc::new(MNIST::new(args, xfrm))
    }
}

impl<T: NumLimits> MNIST<T> {
    pub fn build(root: &str) -> MNISTArgsBuilder {
        MNISTArgsBuilder::default().root(root.into())
    }
    pub fn new(args: MNISTArgs, xfrm: Option<Xfrm>) -> Self {
        if args.download {
            download(&args.root).expect("download failed");
        }
        let processed_path = PathBuf::from(&args.root).join(PROCESSED_FOLDER);
        if !check_exists(&processed_path).expect("Dataset not found, try downloading") {
            panic!("Dataset not found, use download=true to download it");
        }
        let mut v: Vec<TensorKind> = if args.train {
            torch::load(processed_path.join(TRAINING_FILE)).expect("torch load failed")
        } else {
            torch::load(processed_path.join(TEST_FILE)).expect("torch load failed")
        };
        let (data, labels) = (v.remove(0), v.remove(0));

        MNIST {
            root: args.root.clone(),
            train: args.train,
            data: data.into(),
            labels: labels.into(),
            transform: xfrm,
            phantom: PhantomData,
        }
    }
    fn index(&self, idx: usize) -> Sample<u8> {
        let img = self.data.s([idx]);
        let img = if let Some(ref transform) = self.transform {
            transform(&img.into()).into()
        } else {
            img.clone()
        };
        (img, self.labels[idx].clone() as i64)
    }
}

impl<T: NumLimits> DatasetIntf for MNIST<T> {
    type Batch = CollatedSample<T>;
    fn len(&self) -> usize {
        if self.train { 60000 } else { 10000 }
    }
    fn collate(&self, sample: Vec<usize>) -> Self::Batch {
        println!("collate called len: {}", sample.len());
        let v: Vec<Sample<u8>> = sample.into_iter().map(|i| self.index(i)).collect();
        let labels: Vec<i64> = v.iter().map(|&(_, ref t)| *t).collect();
        let imgs: Vec<Tensor<u8>> = v.into_iter().map(|(d, _)| d).collect();
        // XXX should check for case of double
        let img_batch: Tensor<T> = torch::tensor(imgs);
        let label_batch = torch::long_tensor(labels);
        (img_batch, label_batch)
    }
}
<|endoftext|>


<commit_before>use std::thread;
use std::sync::{Arc, Mutex, MutexGuard};

use time::{self, Duration, Timespec};

use connection::{Event, Writer};
use message::Prefix;

#[derive(Clone)]
enum MonitorStatus {
    // Last activity, timespec is when.
    Activity(Timespec),
    // Ping was sent, timespec is when.
    Ping(Timespec),
}

#[derive(Clone)]
enum ConnectionStatus {
    // Connection is alive and well.
    Connected(MonitorStatus),
    // Connection was dropped.
    Disconnected,
    // When the monitor receives a close event, it sets the status
    // to this value to let the background thread know it should stop.
    Quit,
}

#[derive(Clone)]
struct State {
    // Holds the status of the connection.
    status: Arc<Mutex<ConnectionStatus>>,
    // When sending a ping message on irc, you need to specify the server.
    // This holds the current server's name, it's grabbed from the event
    // stream by the feed method.
    server: Arc<Mutex<Option<String>>>,
}

impl State {

    fn new(ts: Timespec) -> State {
        let conn_status = ConnectionStatus::Connected(MonitorStatus::Activity(ts));
        State {
            status: Arc::new(Mutex::new(conn_status)),
            server: Arc::new(Mutex::new(None)),
        }
    }

    // Set the last activity's timestamp.
    fn set_activity(&self, ts: Timespec) {
        *self.status.lock().unwrap() = ConnectionStatus::Connected(MonitorStatus::Activity(ts));
    }

    // Set the status to disconnected.
    fn set_disconnected(&self) {
        *self.status.lock().unwrap() = ConnectionStatus::Disconnected;
    }

    // Set the status to quit.
    fn quit(&self) {
        *self.status.lock().unwrap() = ConnectionStatus::Quit;
    }

    // Get the connection status.
    fn connection_status<'a>(&self) -> MutexGuard<ConnectionStatus> {
        self.status.lock().unwrap()
    }

    // Check if the server variable is set.
    fn has_server(&self) -> bool {
        self.server.lock().unwrap().is_some()
    }

    // Get the server's name.
    fn get_server(&self) -> MutexGuard<Option<String>> {
        self.server.lock().unwrap()
    }

    // Set the server's name.
    fn set_server(&self, name: String) {
        *self.server.lock().unwrap() = Some(name);
    }

    // Unsert the server name. Used when a disconnection occurs.
    fn unset_server(&self) {
        *self.server.lock().unwrap() = None;
    }

}

fn periodic_checker(state: State, handle: Writer, settings: MonitorSettings) {
    loop {
        let mut conn_status = state.connection_status();

        // Clone here, to we can mutate the status inside the match.
        // The type is Copy anyway.
        match conn_status.clone() {
            ConnectionStatus::Connected(ref mon_status) => {
                match *mon_status {
                    // The monitor is in activity mode.
                    // If the timer expires, it will ping the server and set the monitor status
                    // to ping mode.
                    MonitorStatus::Activity(activity_ts) => {
                        let diff = time::get_time() - activity_ts;
                        if diff > settings.activity_timeout {
                            println!("sending a ping");
                            // Make sure we have a server name.
                            match *state.get_server() {
                                Some(ref server) =>  {
                                    // Set the monitor's status to ping mode.
                                    *conn_status = ConnectionStatus::Connected(MonitorStatus::Ping(time::get_time()));
                                    // Send a ping, which should trigger activity is the connection is still alive.
                                    let _ = handle.ping(server);
                                }
                                None => {
                                    // This should really never happen. The server's name should be grabbed
                                    // almost instantly, when the first message is received from the server
                                    // after establishing a new connection.
                                    panic!("Server is None! This scenario is highly unlikely, please report this issue!");
                                }
                            }

                        }
                    }
                    // The monitor is in ping mode, which means it expects a ping response anytime.
                    // If the timer expires, the connection is set to disconnect mode.
                    MonitorStatus::Ping(ping_ts) => {
                        let diff = time::get_time() - ping_ts;
                        println!("checking {}", diff);
                        if diff > settings.ping_timeout {
                            println!("no ping reply");
                            // trigger reconnection process
                            let _ = handle.disconnect();
                        }
                    },
                }
            },
            // Do nothing if the socket is disconnected, for now.
            ConnectionStatus::Disconnected => {},
            ConnectionStatus::Quit => break,
        }

        // Drop the lock to avoid a dead lock.
        drop(conn_status);
        thread::sleep_ms(1000);
    }
}

/// These settings tell the monitor how to behave.
///
/// They allow you to configure the amount of time between the steps.
#[derive(Copy, Clone)]
pub struct MonitorSettings {
    /// Amount of time since the last activity.
    ///
    /// When the amount of time since the last activity gets higher than this value,
    /// tt will trigger a ping request. If there is not a lot of activity, this means
    /// that the server will be pinged everytime this duration expires.
    pub activity_timeout: Duration,
    /// Amount of time to wait for a ping reply.
    ///
    /// When the amount of time since the ping was sent gets higher than this value,
    /// and that no activity occured, assume the connection was dropped and trigger
    /// and a disconnect.
    pub ping_timeout: Duration,
}

/// Default values are provided for the settings.
///
/// They are:
///
/// `activity_timeout` = 60 seconds
///
/// `ping_timeout` = 15 seconds
impl Default for MonitorSettings {

    fn default() -> MonitorSettings {
        MonitorSettings {
            activity_timeout: Duration::seconds(60),
            ping_timeout: Duration::seconds(15),
        }
    }

}

/// This struct monitors a connection's activity.
///
/// It works in a few simple, steps.
/// First, it monitors activity via the feed method, saving a timestamp for relevant events.
/// If no activity is detected over the given period of time, it will send a ping request to
/// the IRC server. If no ping  reply is received, it will trigger a disconnect which may lead
/// to reconnection depending on the reconnection settings.
///
/// The amount of time to wait, without activity, before sending a ping request, and the amount
/// of time to wait for the ping reply can be configured via the `MonitorSettings` struct.
pub struct ActivityMonitor {
    state: State,
}

impl ActivityMonitor {

    /// Create a new ActivityMonitor.
    ///
    /// The handle to a Writer allows the monitor to notify the connection of disconnects.
    pub fn new(handle: &Writer, settings: MonitorSettings) -> ActivityMonitor {
        let state =  State::new(time::get_time());

        let state_clone = state.clone();
        let handle_clone = handle.clone();

        thread::spawn(move || {
            periodic_checker(state_clone, handle_clone, settings);
        });

        ActivityMonitor {
            state: state,
        }
    }

    /// Give an event received from the connection to the monitor.
    ///
    /// The monitor will process it accordingly. It an Event::Closed event
    /// is received, it will shutdown all of its activities.
    pub fn feed(&self, event: &Event) {
        match *event {
            Event::Closed(_) => {
                self.state.quit();
            }
            Event::Disconnected => {
                self.state.set_disconnected();
                self.state.unset_server();
            }
            Event::Reconnected => {
                self.state.set_activity(time::get_time());
            }
            Event::Message(ref msg) => {
                self.state.set_activity(time::get_time());
                if let Some(ref prefix) = msg.prefix {
                    match *prefix {
                        Prefix::Server(ref name) => {
                            if !self.state.has_server() {
                                self.state.set_server(name.clone());
                            }
                        }
                        _ => {}
                    }
                }
            }
            // Other events are irrelevant.
            _ => {}
        }
    }

}

/// Drop stops the background thread and clears the monitor's resources.
///
/// If you can the activity monitor to cease its activites, you can simply drop it.
/// It will not affect the connection on which the activity monitor operates.
impl Drop for ActivityMonitor {

    fn drop(&mut self) {
        self.state.quit();
    }

}
<commit_msg>Remove println! statements from ActivityMonitor<commit_after>use std::thread;
use std::sync::{Arc, Mutex, MutexGuard};

use time::{self, Duration, Timespec};

use connection::{Event, Writer};
use message::Prefix;

#[derive(Clone)]
enum MonitorStatus {
    // Last activity, timespec is when.
    Activity(Timespec),
    // Ping was sent, timespec is when.
    Ping(Timespec),
}

#[derive(Clone)]
enum ConnectionStatus {
    // Connection is alive and well.
    Connected(MonitorStatus),
    // Connection was dropped.
    Disconnected,
    // When the monitor receives a close event, it sets the status
    // to this value to let the background thread know it should stop.
    Quit,
}

#[derive(Clone)]
struct State {
    // Holds the status of the connection.
    status: Arc<Mutex<ConnectionStatus>>,
    // When sending a ping message on irc, you need to specify the server.
    // This holds the current server's name, it's grabbed from the event
    // stream by the feed method.
    server: Arc<Mutex<Option<String>>>,
}

impl State {

    fn new(ts: Timespec) -> State {
        let conn_status = ConnectionStatus::Connected(MonitorStatus::Activity(ts));
        State {
            status: Arc::new(Mutex::new(conn_status)),
            server: Arc::new(Mutex::new(None)),
        }
    }

    // Set the last activity's timestamp.
    fn set_activity(&self, ts: Timespec) {
        *self.status.lock().unwrap() = ConnectionStatus::Connected(MonitorStatus::Activity(ts));
    }

    // Set the status to disconnected.
    fn set_disconnected(&self) {
        *self.status.lock().unwrap() = ConnectionStatus::Disconnected;
    }

    // Set the status to quit.
    fn quit(&self) {
        *self.status.lock().unwrap() = ConnectionStatus::Quit;
    }

    // Get the connection status.
    fn connection_status<'a>(&self) -> MutexGuard<ConnectionStatus> {
        self.status.lock().unwrap()
    }

    // Check if the server variable is set.
    fn has_server(&self) -> bool {
        self.server.lock().unwrap().is_some()
    }

    // Get the server's name.
    fn get_server(&self) -> MutexGuard<Option<String>> {
        self.server.lock().unwrap()
    }

    // Set the server's name.
    fn set_server(&self, name: String) {
        *self.server.lock().unwrap() = Some(name);
    }

    // Unsert the server name. Used when a disconnection occurs.
    fn unset_server(&self) {
        *self.server.lock().unwrap() = None;
    }

}

fn periodic_checker(state: State, handle: Writer, settings: MonitorSettings) {
    loop {
        let mut conn_status = state.connection_status();

        // Clone here, to we can mutate the status inside the match.
        // The type is Copy anyway.
        match conn_status.clone() {
            ConnectionStatus::Connected(ref mon_status) => {
                match *mon_status {
                    // The monitor is in activity mode.
                    // If the timer expires, it will ping the server and set the monitor status
                    // to ping mode.
                    MonitorStatus::Activity(activity_ts) => {
                        let diff = time::get_time() - activity_ts;
                        if diff > settings.activity_timeout {
                            // Make sure we have a server name.
                            match *state.get_server() {
                                Some(ref server) =>  {
                                    // Set the monitor's status to ping mode.
                                    *conn_status = ConnectionStatus::Connected(MonitorStatus::Ping(time::get_time()));
                                    // Send a ping, which should trigger activity is the connection is still alive.
                                    let _ = handle.ping(server);
                                }
                                None => {
                                    // This should really never happen. The server's name should be grabbed
                                    // almost instantly, when the first message is received from the server
                                    // after establishing a new connection.
                                    panic!("Server is None! This scenario is highly unlikely, please report this issue!");
                                }
                            }

                        }
                    }
                    // The monitor is in ping mode, which means it expects a ping response anytime.
                    // If the timer expires, the connection is set to disconnect mode.
                    MonitorStatus::Ping(ping_ts) => {
                        let diff = time::get_time() - ping_ts;
                        if diff > settings.ping_timeout {
                            // trigger reconnection process
                            let _ = handle.disconnect();
                        }
                    },
                }
            },
            // Do nothing if the socket is disconnected, for now.
            ConnectionStatus::Disconnected => {},
            ConnectionStatus::Quit => break,
        }

        // Drop the lock to avoid a dead lock.
        drop(conn_status);
        thread::sleep_ms(1000);
    }
}

/// These settings tell the monitor how to behave.
///
/// They allow you to configure the amount of time between the steps.
#[derive(Copy, Clone)]
pub struct MonitorSettings {
    /// Amount of time since the last activity.
    ///
    /// When the amount of time since the last activity gets higher than this value,
    /// tt will trigger a ping request. If there is not a lot of activity, this means
    /// that the server will be pinged everytime this duration expires.
    pub activity_timeout: Duration,
    /// Amount of time to wait for a ping reply.
    ///
    /// When the amount of time since the ping was sent gets higher than this value,
    /// and that no activity occured, assume the connection was dropped and trigger
    /// and a disconnect.
    pub ping_timeout: Duration,
}

/// Default values are provided for the settings.
///
/// They are:
///
/// `activity_timeout` = 60 seconds
///
/// `ping_timeout` = 15 seconds
impl Default for MonitorSettings {

    fn default() -> MonitorSettings {
        MonitorSettings {
            activity_timeout: Duration::seconds(60),
            ping_timeout: Duration::seconds(15),
        }
    }

}

/// This struct monitors a connection's activity.
///
/// It works in a few simple, steps.
/// First, it monitors activity via the feed method, saving a timestamp for relevant events.
/// If no activity is detected over the given period of time, it will send a ping request to
/// the IRC server. If no ping  reply is received, it will trigger a disconnect which may lead
/// to reconnection depending on the reconnection settings.
///
/// The amount of time to wait, without activity, before sending a ping request, and the amount
/// of time to wait for the ping reply can be configured via the `MonitorSettings` struct.
pub struct ActivityMonitor {
    state: State,
}

impl ActivityMonitor {

    /// Create a new ActivityMonitor.
    ///
    /// The handle to a Writer allows the monitor to notify the connection of disconnects.
    pub fn new(handle: &Writer, settings: MonitorSettings) -> ActivityMonitor {
        let state =  State::new(time::get_time());

        let state_clone = state.clone();
        let handle_clone = handle.clone();

        thread::spawn(move || {
            periodic_checker(state_clone, handle_clone, settings);
        });

        ActivityMonitor {
            state: state,
        }
    }

    /// Give an event received from the connection to the monitor.
    ///
    /// The monitor will process it accordingly. It an Event::Closed event
    /// is received, it will shutdown all of its activities.
    pub fn feed(&self, event: &Event) {
        match *event {
            Event::Closed(_) => {
                self.state.quit();
            }
            Event::Disconnected => {
                self.state.set_disconnected();
                self.state.unset_server();
            }
            Event::Reconnected => {
                self.state.set_activity(time::get_time());
            }
            Event::Message(ref msg) => {
                self.state.set_activity(time::get_time());
                if let Some(ref prefix) = msg.prefix {
                    match *prefix {
                        Prefix::Server(ref name) => {
                            if !self.state.has_server() {
                                self.state.set_server(name.clone());
                            }
                        }
                        _ => {}
                    }
                }
            }
            // Other events are irrelevant.
            _ => {}
        }
    }

}

/// Drop stops the background thread and clears the monitor's resources.
///
/// If you can the activity monitor to cease its activites, you can simply drop it.
/// It will not affect the connection on which the activity monitor operates.
impl Drop for ActivityMonitor {

    fn drop(&mut self) {
        self.state.quit();
    }

}
<|endoftext|>


<commit_before>mod data;

pub use self::data::AirportType;

use self::data::DataType;
use crate::path::FilePath;
use anyhow::{anyhow, Result};
use chrono::{Duration, NaiveDate, Utc};
use serde_derive::Serialize;
use std::borrow::Cow;
use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{self, BufWriter};
use std::path::{Path, PathBuf};

const BASE_URL: &str = "https://ourairports.com/data";

fn validated_dir() -> Result<PathBuf> {
    FilePath::LocalData.validated_subdir("airport_data/")
}

pub fn ensure_updated() -> Result<()> {
    let base_dir = validated_dir()?;
    let mut last_update = LastUpdate::load(&base_dir);

    if !last_update.needs_update() {
        return Ok(());
    }

    download_latest(&base_dir)?;
    last_update.set_to_today()?;

    Ok(())
}

fn download_latest<P>(dir: P) -> Result<()>
where
    P: AsRef<Path>,
{
    println!("updating airport data");

    for dtype in &DataType::ALL {
        download_file(dtype.filename(), &dir)?;
    }

    println!("finished updating airport data");
    Ok(())
}

fn download_file<S, P>(name: S, dir: P) -> Result<()>
where
    S: AsRef<str>,
    P: AsRef<Path>,
{
    let name = name.as_ref();
    let dir = dir.as_ref();

    println!(".. downloading {}", name);

    let path = dir.join(name);
    let backup = Backup::create(&path)?;

    let resp = attohttpc::get(format!("{}/{}", BASE_URL, name))
        .timeout(Duration::seconds(15).to_std().unwrap())
        .send()?;

    if !resp.is_success() {
        backup.restore()?;

        return Err(anyhow!(
            "received code {} while downloading {}",
            resp.status(),
            name
        ));
    }

    let file = match File::create(path) {
        Ok(file) => file,
        Err(err) => {
            backup.restore()?;
            return Err(err.into());
        }
    };

    let (_, _, mut reader) = resp.split();
    let mut writer = BufWriter::new(file);

    match io::copy(&mut reader, &mut writer) {
        Ok(_) => Ok(()),
        Err(err) => {
            backup.restore()?;
            Err(err.into())
        }
    }
}

struct LastUpdate(PathBuf);

impl LastUpdate {
    const DATE_FORMAT: &'static str = "%Y-%m-%d";

    fn load<P>(base_dir: P) -> Self
    where
        P: Into<PathBuf>,
    {
        let mut path = base_dir.into();
        path.push(".last_updated");

        Self(path)
    }

    fn needs_update(&self) -> bool {
        if !self.path().exists() {
            return true;
        }

        let updated_at = fs::read_to_string(self.path())
            .ok()
            .and_then(|updated_at| NaiveDate::parse_from_str(&updated_at, Self::DATE_FORMAT).ok());

        let updated_at = match updated_at {
            Some(updated_at) => updated_at,
            None => return true,
        };

        let today = Utc::now().naive_utc().date();

        if today - updated_at >= Duration::weeks(2) {
            return true;
        }

        false
    }

    fn set_to_today(&mut self) -> Result<()> {
        let today = Utc::now().naive_utc().date();
        fs::write(self.path(), today.format(Self::DATE_FORMAT).to_string())?;
        Ok(())
    }

    #[inline(always)]
    fn path(&self) -> &Path {
        &self.0
    }
}

enum Backup {
    Exists(PathBuf),
    DoesntExist,
}

impl Backup {
    fn create<P>(file: P) -> Result<Self>
    where
        P: AsRef<Path>,
    {
        let file = file.as_ref();

        if !file.exists() {
            return Ok(Self::DoesntExist);
        }

        let extension = match file.extension() {
            Some(ext) => format!("{}.bak", ext.to_string_lossy()).into(),
            None => Cow::Borrowed(".bak"),
        };

        let mut backup_path = PathBuf::from(file);
        backup_path.set_extension(extension.as_ref());

        fs::rename(file, &backup_path)?;

        Ok(Self::Exists(backup_path))
    }

    fn restore(self) -> Result<()> {
        match self {
            Self::Exists(path) => {
                if !path.exists() {
                    return Ok(());
                }

                let mut original_path = path.clone();
                original_path.set_extension("");

                if original_path.exists() {
                    fs::remove_file(&original_path)?;
                }

                fs::rename(path, original_path)?;
                Ok(())
            }
            Self::DoesntExist => Ok(()),
        }
    }
}

#[derive(Debug, Serialize)]
pub struct Airport {
    pub icao: String,
    #[serde(skip_serializing)]
    pub class: AirportType,
    pub position: Position,
    pub runways: Vec<Runway>,
    pub frequencies: HashMap<FrequencyType, String>,
    #[serde(skip_serializing)]
    pub country_name: String,
}

impl Airport {
    pub fn load_all() -> Result<Vec<Self>> {
        let dir = validated_dir()?;

        let airports = data::Airport::from_dir(&dir)?;
        let mut runways = data::Runway::from_dir(&dir)?;
        let mut frequencies = data::Frequency::from_dir(&dir)?;
        let countries = data::Country::from_dir(&dir)?;

        let mut results = Vec::with_capacity(airports.len());

        for airport in airports {
            let runways = match runways.remove(&airport.id) {
                Some(runways) => runways,
                None => continue,
            };

            let frequencies = match frequencies.remove(&airport.id) {
                Some(frequencies) => frequencies,
                None => HashMap::new(),
            };

            let country = match countries.iter().find(|c| c.code == airport.country_code) {
                Some(country) => country,
                None => continue,
            };

            let result = Self {
                icao: airport.icao,
                class: airport.class,
                position: Position::new(airport.lat_deg, airport.lon_deg),
                runways: runways.into_iter().filter_map(Runway::from_data).collect(),
                frequencies: frequencies
                    .into_iter()
                    .map(|(freq_type, freq)| (freq_type.into(), freq.mhz))
                    .collect(),
                country_name: country.name.clone(),
            };

            results.push(result);
        }

        results.shrink_to_fit();
        results.sort_unstable_by(|x, y| x.icao.cmp(&y.icao));

        Ok(results)
    }
}

impl PartialEq for Airport {
    fn eq(&self, other: &Self) -> bool {
        self.icao == other.icao
    }
}

#[derive(Copy, Clone, Debug, Serialize)]
pub struct Position {
    #[serde(rename = "latitudeDeg")]
    pub latitude_deg: f32,
    #[serde(rename = "longitudeDeg")]
    pub longitude_deg: f32,
}

impl Position {
    #[inline(always)]
    pub fn new(latitude_deg: f32, longitude_deg: f32) -> Self {
        Self {
            latitude_deg,
            longitude_deg,
        }
    }
}

#[derive(Debug, Serialize)]
pub struct Runway {
    #[serde(rename = "lengthFT")]
    pub length_ft: Option<u32>,
    #[serde(rename = "widthFT")]
    pub width_ft: Option<u32>,
    #[serde(rename = "heMarker")]
    pub he_marker: Option<RunwayMarker>,
    #[serde(rename = "leMarker")]
    pub le_marker: Option<RunwayMarker>,
}

impl Runway {
    fn from_data(source: data::Runway) -> Option<Self> {
        let he_marker = match (source.le_ident, source.le_lat_deg, source.le_lon_deg) {
            (Some(name), Some(lat), Some(lon)) => Some(RunwayMarker::new(name, lat, lon)),
            _ => None,
        };

        let le_marker = match (source.he_ident, source.he_lat_deg, source.he_lon_deg) {
            (Some(name), Some(lat), Some(lon)) => Some(RunwayMarker::new(name, lat, lon)),
            _ => None,
        };

        Some(Self {
            length_ft: source.length_ft,
            width_ft: source.width_ft,
            he_marker,
            le_marker,
        })
    }
}

#[derive(Debug, Serialize)]
pub struct RunwayMarker {
    pub name: String,
    pub position: Position,
}

impl RunwayMarker {
    #[inline(always)]
    pub fn new(name: String, lat_deg: f32, lon_deg: f32) -> Self {
        Self {
            name,
            position: Position::new(lat_deg, lon_deg),
        }
    }
}

#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum FrequencyType {
    Atis,
    Arrival,
    Departure,
    #[serde(rename = "arrivalDeparture")]
    ArrivalDeparture,
    Ground,
    Tower,
    Unicom,
}

// TODO: return None when the source type is the Other variant
impl From<data::FrequencyType> for FrequencyType {
    fn from(source: data::FrequencyType) -> Self {
        use data::FrequencyType as DFT;

        match source {
            DFT::Atis => Self::Atis,
            DFT::Arrival => Self::Arrival,
            DFT::Departure => Self::Departure,
            DFT::ArrivalDeparture => Self::ArrivalDeparture,
            DFT::Ground => Self::Ground,
            DFT::Tower => Self::Tower,
            DFT::Unicom | DFT::Other => Self::Unicom,
        }
    }
}
<commit_msg>backend: Limit max ICAO length to 4 characters<commit_after>mod data;

pub use self::data::AirportType;

use self::data::DataType;
use crate::path::FilePath;
use anyhow::{anyhow, Result};
use chrono::{Duration, NaiveDate, Utc};
use serde_derive::Serialize;
use std::borrow::Cow;
use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{self, BufWriter};
use std::path::{Path, PathBuf};

const BASE_URL: &str = "https://ourairports.com/data";

fn validated_dir() -> Result<PathBuf> {
    FilePath::LocalData.validated_subdir("airport_data/")
}

pub fn ensure_updated() -> Result<()> {
    let base_dir = validated_dir()?;
    let mut last_update = LastUpdate::load(&base_dir);

    if !last_update.needs_update() {
        return Ok(());
    }

    download_latest(&base_dir)?;
    last_update.set_to_today()?;

    Ok(())
}

fn download_latest<P>(dir: P) -> Result<()>
where
    P: AsRef<Path>,
{
    println!("updating airport data");

    for dtype in &DataType::ALL {
        download_file(dtype.filename(), &dir)?;
    }

    println!("finished updating airport data");
    Ok(())
}

fn download_file<S, P>(name: S, dir: P) -> Result<()>
where
    S: AsRef<str>,
    P: AsRef<Path>,
{
    let name = name.as_ref();
    let dir = dir.as_ref();

    println!(".. downloading {}", name);

    let path = dir.join(name);
    let backup = Backup::create(&path)?;

    let resp = attohttpc::get(format!("{}/{}", BASE_URL, name))
        .timeout(Duration::seconds(15).to_std().unwrap())
        .send()?;

    if !resp.is_success() {
        backup.restore()?;

        return Err(anyhow!(
            "received code {} while downloading {}",
            resp.status(),
            name
        ));
    }

    let file = match File::create(path) {
        Ok(file) => file,
        Err(err) => {
            backup.restore()?;
            return Err(err.into());
        }
    };

    let (_, _, mut reader) = resp.split();
    let mut writer = BufWriter::new(file);

    match io::copy(&mut reader, &mut writer) {
        Ok(_) => Ok(()),
        Err(err) => {
            backup.restore()?;
            Err(err.into())
        }
    }
}

struct LastUpdate(PathBuf);

impl LastUpdate {
    const DATE_FORMAT: &'static str = "%Y-%m-%d";

    fn load<P>(base_dir: P) -> Self
    where
        P: Into<PathBuf>,
    {
        let mut path = base_dir.into();
        path.push(".last_updated");

        Self(path)
    }

    fn needs_update(&self) -> bool {
        if !self.path().exists() {
            return true;
        }

        let updated_at = fs::read_to_string(self.path())
            .ok()
            .and_then(|updated_at| NaiveDate::parse_from_str(&updated_at, Self::DATE_FORMAT).ok());

        let updated_at = match updated_at {
            Some(updated_at) => updated_at,
            None => return true,
        };

        let today = Utc::now().naive_utc().date();

        if today - updated_at >= Duration::weeks(2) {
            return true;
        }

        false
    }

    fn set_to_today(&mut self) -> Result<()> {
        let today = Utc::now().naive_utc().date();
        fs::write(self.path(), today.format(Self::DATE_FORMAT).to_string())?;
        Ok(())
    }

    #[inline(always)]
    fn path(&self) -> &Path {
        &self.0
    }
}

enum Backup {
    Exists(PathBuf),
    DoesntExist,
}

impl Backup {
    fn create<P>(file: P) -> Result<Self>
    where
        P: AsRef<Path>,
    {
        let file = file.as_ref();

        if !file.exists() {
            return Ok(Self::DoesntExist);
        }

        let extension = match file.extension() {
            Some(ext) => format!("{}.bak", ext.to_string_lossy()).into(),
            None => Cow::Borrowed(".bak"),
        };

        let mut backup_path = PathBuf::from(file);
        backup_path.set_extension(extension.as_ref());

        fs::rename(file, &backup_path)?;

        Ok(Self::Exists(backup_path))
    }

    fn restore(self) -> Result<()> {
        match self {
            Self::Exists(path) => {
                if !path.exists() {
                    return Ok(());
                }

                let mut original_path = path.clone();
                original_path.set_extension("");

                if original_path.exists() {
                    fs::remove_file(&original_path)?;
                }

                fs::rename(path, original_path)?;
                Ok(())
            }
            Self::DoesntExist => Ok(()),
        }
    }
}

#[derive(Debug, Serialize)]
pub struct Airport {
    pub icao: String,
    #[serde(skip_serializing)]
    pub class: AirportType,
    pub position: Position,
    pub runways: Vec<Runway>,
    pub frequencies: HashMap<FrequencyType, String>,
    #[serde(skip_serializing)]
    pub country_name: String,
}

impl Airport {
    const ICAO_MAX_LEN: usize = 4;

    pub fn load_all() -> Result<Vec<Self>> {
        let dir = validated_dir()?;

        let airports = data::Airport::from_dir(&dir)?;
        let mut runways = data::Runway::from_dir(&dir)?;
        let mut frequencies = data::Frequency::from_dir(&dir)?;
        let countries = data::Country::from_dir(&dir)?;

        let mut results = Vec::with_capacity(airports.len());

        for airport in airports {
            if airport.icao.len() > Self::ICAO_MAX_LEN {
                continue;
            }

            let runways = match runways.remove(&airport.id) {
                Some(runways) => runways,
                None => continue,
            };

            let frequencies = match frequencies.remove(&airport.id) {
                Some(frequencies) => frequencies,
                None => HashMap::new(),
            };

            let country = match countries.iter().find(|c| c.code == airport.country_code) {
                Some(country) => country,
                None => continue,
            };

            let result = Self {
                icao: airport.icao,
                class: airport.class,
                position: Position::new(airport.lat_deg, airport.lon_deg),
                runways: runways.into_iter().filter_map(Runway::from_data).collect(),
                frequencies: frequencies
                    .into_iter()
                    .map(|(freq_type, freq)| (freq_type.into(), freq.mhz))
                    .collect(),
                country_name: country.name.clone(),
            };

            results.push(result);
        }

        results.shrink_to_fit();
        results.sort_unstable_by(|x, y| x.icao.cmp(&y.icao));

        Ok(results)
    }
}

impl PartialEq for Airport {
    fn eq(&self, other: &Self) -> bool {
        self.icao == other.icao
    }
}

#[derive(Copy, Clone, Debug, Serialize)]
pub struct Position {
    #[serde(rename = "latitudeDeg")]
    pub latitude_deg: f32,
    #[serde(rename = "longitudeDeg")]
    pub longitude_deg: f32,
}

impl Position {
    #[inline(always)]
    pub fn new(latitude_deg: f32, longitude_deg: f32) -> Self {
        Self {
            latitude_deg,
            longitude_deg,
        }
    }
}

#[derive(Debug, Serialize)]
pub struct Runway {
    #[serde(rename = "lengthFT")]
    pub length_ft: Option<u32>,
    #[serde(rename = "widthFT")]
    pub width_ft: Option<u32>,
    #[serde(rename = "heMarker")]
    pub he_marker: Option<RunwayMarker>,
    #[serde(rename = "leMarker")]
    pub le_marker: Option<RunwayMarker>,
}

impl Runway {
    fn from_data(source: data::Runway) -> Option<Self> {
        let he_marker = match (source.le_ident, source.le_lat_deg, source.le_lon_deg) {
            (Some(name), Some(lat), Some(lon)) => Some(RunwayMarker::new(name, lat, lon)),
            _ => None,
        };

        let le_marker = match (source.he_ident, source.he_lat_deg, source.he_lon_deg) {
            (Some(name), Some(lat), Some(lon)) => Some(RunwayMarker::new(name, lat, lon)),
            _ => None,
        };

        Some(Self {
            length_ft: source.length_ft,
            width_ft: source.width_ft,
            he_marker,
            le_marker,
        })
    }
}

#[derive(Debug, Serialize)]
pub struct RunwayMarker {
    pub name: String,
    pub position: Position,
}

impl RunwayMarker {
    #[inline(always)]
    pub fn new(name: String, lat_deg: f32, lon_deg: f32) -> Self {
        Self {
            name,
            position: Position::new(lat_deg, lon_deg),
        }
    }
}

#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum FrequencyType {
    Atis,
    Arrival,
    Departure,
    #[serde(rename = "arrivalDeparture")]
    ArrivalDeparture,
    Ground,
    Tower,
    Unicom,
}

// TODO: return None when the source type is the Other variant
impl From<data::FrequencyType> for FrequencyType {
    fn from(source: data::FrequencyType) -> Self {
        use data::FrequencyType as DFT;

        match source {
            DFT::Atis => Self::Atis,
            DFT::Arrival => Self::Arrival,
            DFT::Departure => Self::Departure,
            DFT::ArrivalDeparture => Self::ArrivalDeparture,
            DFT::Ground => Self::Ground,
            DFT::Tower => Self::Tower,
            DFT::Unicom | DFT::Other => Self::Unicom,
        }
    }
}
<|endoftext|>


<commit_before>extern crate sdl2;

use sdl2::audio::{AudioCallback, AudioSpecDesired,AudioSpecWAV,AudioCVT};
use std::time::Duration;
use std::borrow::Cow;
use std::path::{PathBuf, Path};

// NOTE: You probably want to investigate the
// mixer feature for real use cases.

struct Sound {
    data: Vec<u8>,
    volume: f32,
    pos: usize,
}

impl AudioCallback for Sound {
    type Channel = u8;

    fn callback(&mut self, out: &mut [u8]) {
        for dst in out.iter_mut() {
            *dst = (*self.data.get(self.pos).unwrap_or(&0) as f32 * self.volume) as u8;
            self.pos += 1;
        }
    }
}

fn main() -> Result<(), String> {
    let wav_file : Cow<'static, Path> = match std::env::args().nth(1) {
        None => Cow::from(Path::new("./assets/sine.wav")),
        Some(s) => Cow::from(PathBuf::from(s))
    };
    let sdl_context = sdl2::init()?;
    let audio_subsystem = sdl_context.audio()?;

    let desired_spec = AudioSpecDesired {
        freq: Some(44_100),
        channels: Some(1), // mono
        samples: None      // default
    };

    let device = audio_subsystem.open_playback(None, &desired_spec, |spec| {
        let wav = AudioSpecWAV::load_wav(wav_file)
            .expect("Could not load test WAV file");

        let cvt = AudioCVT::new(
                wav.format, wav.channels, wav.freq,
                spec.format, spec.channels, spec.freq)
            .expect("Could not convert WAV file");

        let data = cvt.convert(wav.buffer().to_vec());

        // initialize the audio callback
        Sound {
            data: data,
            volume: 0.25,
            pos: 0,
        }
    })?;

    // Start playback
    device.resume();

    // Play for a second
    std::thread::sleep(Duration::from_millis(1_000));

    // Device is automatically closed when dropped

    Ok(())
}
<commit_msg>Fix clicking and volume control in the audio-wav example<commit_after>extern crate sdl2;

use sdl2::audio::{AudioCallback, AudioSpecDesired,AudioSpecWAV,AudioCVT};
use std::time::Duration;
use std::borrow::Cow;
use std::path::{PathBuf, Path};

// NOTE: You probably want to investigate the
// mixer feature for real use cases.

struct Sound {
    data: Vec<u8>,
    volume: f32,
    pos: usize,
}

impl AudioCallback for Sound {
    type Channel = u8;

    fn callback(&mut self, out: &mut [u8]) {
        for dst in out.iter_mut() {
            // With channel type u8 the "silence" value is 128 (middle of the 0-2^8 range) so we need
            // to both fill in the silence and scale the wav data accordingly. Filling the silence
            // once the wav is finished is trivial, applying the volume is more tricky. We need to:
            // * Change the range of the values from [0, 255] to [-128, 127] so we can multiply
            // * Apply the volume by multiplying, this gives us range [-128*volume, 127*volume]
            // * Move the resulting range to a range centered around the value 128, the final range
            //   is [128 - 128*volume, 128 + 127*volume] – scaled and correctly positioned
            //
            // Using value 0 instead of 128 would result in clicking. Scaling by simply multiplying
            // would not give correct results.
            let pre_scale = *self.data.get(self.pos).unwrap_or(&128);
            let scaled_signed_float = (pre_scale as f32 - 128.0) * self.volume;
            let scaled = (scaled_signed_float + 128.0) as u8;
            *dst = scaled;
            self.pos += 1;
        }
    }
}

fn main() -> Result<(), String> {
    let wav_file : Cow<'static, Path> = match std::env::args().nth(1) {
        None => Cow::from(Path::new("./assets/sine.wav")),
        Some(s) => Cow::from(PathBuf::from(s))
    };
    let sdl_context = sdl2::init()?;
    let audio_subsystem = sdl_context.audio()?;

    let desired_spec = AudioSpecDesired {
        freq: Some(44_100),
        channels: Some(1), // mono
        samples: None      // default
    };

    let device = audio_subsystem.open_playback(None, &desired_spec, |spec| {
        let wav = AudioSpecWAV::load_wav(wav_file)
            .expect("Could not load test WAV file");

        let cvt = AudioCVT::new(
                wav.format, wav.channels, wav.freq,
                spec.format, spec.channels, spec.freq)
            .expect("Could not convert WAV file");

        let data = cvt.convert(wav.buffer().to_vec());

        // initialize the audio callback
        Sound {
            data: data,
            volume: 0.25,
            pos: 0,
        }
    })?;

    // Start playback
    device.resume();

    // Play for a second
    std::thread::sleep(Duration::from_millis(1_000));

    // Device is automatically closed when dropped

    Ok(())
}
<|endoftext|>