<commit_before>use mio::*;
use std::time::Duration;

#[test]
fn smoke() {
    let poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(128);

    let (_r, set) = Registration::new(&poll, Token(0), Ready::readable(), PollOpt::edge());

    let n = poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();
    assert_eq!(n, 0);

    set.set_readiness(Ready::readable()).unwrap();

    let n = poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();
    assert_eq!(n, 1);

    assert_eq!(events.get(0).unwrap().token(), Token(0));
}

#[test]
fn set_readiness_before_register() {
    use std::sync::{Arc, Barrier};
    use std::thread;

    let poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(128);

    for _ in 0..5_000 {
        let (r, set) = Registration::new2();

        let b1 = Arc::new(Barrier::new(2));
        let b2 = b1.clone();

        let th = thread::spawn(move || {
            b2.wait();
            set.set_readiness(Ready::readable()).unwrap();
        });

        b1.wait();

        poll.register(&r, Token(123), Ready::readable(), PollOpt::edge()).unwrap();

        loop {
            let n = poll.poll(&mut events, None).unwrap();

            if n == 0 {
                continue;
            }

            assert_eq!(n, 1);
            assert_eq!(events.get(0).unwrap().token(), Token(123));
            break;
        }

        th.join().unwrap();
    }
}

#[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
mod stress {
    use mio::*;
    use std::time::Duration;

    #[test]
    fn single_threaded_poll() {
        use std::sync::Arc;
        use std::sync::atomic::AtomicUsize;
        use std::sync::atomic::Ordering::{Acquire, Release};
        use std::thread;

        const NUM_ATTEMPTS: usize = 30;
        const NUM_ITERS: usize = 500;
        const NUM_THREADS: usize = 4;
        const NUM_REGISTRATIONS: usize = 128;

        for _ in 0..NUM_ATTEMPTS {
            let poll = Poll::new().unwrap();
            let mut events = Events::with_capacity(128);

            let registrations: Vec<_> = (0..NUM_REGISTRATIONS).map(|i| {
                Registration::new(&poll, Token(i), Ready::readable(), PollOpt::edge())
            }).collect();

            let mut ready: Vec<_> = (0..NUM_REGISTRATIONS).map(|_| Ready::none()).collect();

            let remaining = Arc::new(AtomicUsize::new(NUM_THREADS));

            for _ in 0..NUM_THREADS {
                let remaining = remaining.clone();

                let set_readiness: Vec<SetReadiness> =
                    registrations.iter().map(|r| r.1.clone()).collect();

                thread::spawn(move || {
                    for _ in 0..NUM_ITERS {
                        for i in 0..NUM_REGISTRATIONS {
                            set_readiness[i].set_readiness(Ready::readable()).unwrap();
                            set_readiness[i].set_readiness(Ready::none()).unwrap();
                            set_readiness[i].set_readiness(Ready::writable()).unwrap();
                            set_readiness[i].set_readiness(Ready::readable() | Ready::writable()).unwrap();
                            set_readiness[i].set_readiness(Ready::none()).unwrap();
                        }
                    }

                    for i in 0..NUM_REGISTRATIONS {
                        set_readiness[i].set_readiness(Ready::readable()).unwrap();
                    }

                    remaining.fetch_sub(1, Release);
                });
            }

            while remaining.load(Acquire) > 0 {
                // Set interest
                for (i, &(ref r, _)) in registrations.iter().enumerate() {
                    r.update(&poll, Token(i), Ready::writable(), PollOpt::edge()).unwrap();
                }

                poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();

                for event in &events {
                    ready[event.token().0] = event.kind();
                }

                // Update registration
                // Set interest
                for (i, &(ref r, _)) in registrations.iter().enumerate() {
                    r.update(&poll, Token(i), Ready::readable(), PollOpt::edge()).unwrap();
                }
            }

            // One final poll
            poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();

            for event in &events {
                ready[event.token().0] = event.kind();
            }

            // Everything should be flagged as readable
            for ready in ready {
                assert_eq!(ready, Ready::readable());
            }
        }
    }

    #[test]
    fn multi_threaded_poll() {
        use std::sync::{Arc, Barrier};
        use std::sync::atomic::{AtomicUsize};
        use std::sync::atomic::Ordering::{Relaxed, SeqCst};
        use std::thread;

        const ENTRIES: usize = 10_000;
        const PER_ENTRY: usize = 16;
        const THREADS: usize = 4;
        const NUM: usize = ENTRIES * PER_ENTRY;

        struct Entry {
            #[allow(dead_code)]
            registration: Registration,
            set_readiness: SetReadiness,
            num: AtomicUsize,
        }

        impl Entry {
            fn fire(&self) {
                self.set_readiness.set_readiness(Ready::readable()).unwrap();
            }
        }

        let poll = Arc::new(Poll::new().unwrap());
        let mut entries = vec![];

        // Create entries
        for i in 0..ENTRIES {
            let (registration, set_readiness) =
                Registration::new(&poll, Token(i), Ready::readable(), PollOpt::edge());

            entries.push(Entry {
                registration: registration,
                set_readiness: set_readiness,
                num: AtomicUsize::new(0),
            });
        }

        let total = Arc::new(AtomicUsize::new(0));
        let entries = Arc::new(entries);
        let barrier = Arc::new(Barrier::new(THREADS));

        let mut threads = vec![];

        for th in 0..THREADS {
            let poll = poll.clone();
            let total = total.clone();
            let entries = entries.clone();
            let barrier = barrier.clone();

            threads.push(thread::spawn(move || {
                let mut events = Events::with_capacity(128);

                barrier.wait();

                // Prime all the registrations
                let mut i = th;
                while i < ENTRIES {
                    entries[i].fire();
                    i += THREADS;
                }

                let mut n = 0;

                while total.load(SeqCst) < NUM {
                    // A poll timeout is necessary here because there may be more
                    // than one threads blocked in `poll` when the final wakeup
                    // notification arrives (and only notifies one thread).
                    n += poll.poll(&mut events, Some(Duration::from_millis(100))).unwrap();

                    let mut num_this_tick = 0;

                    for event in &events {
                        let e = &entries[event.token().0];

                        let mut num = e.num.load(Relaxed);

                        loop {
                            if num < PER_ENTRY {
                                let actual = e.num.compare_and_swap(num, num + 1, Relaxed);

                                if actual == num {
                                    num_this_tick += 1;
                                    e.fire();
                                    break;
                                }

                                num = actual;
                            } else {
                                break;
                            }
                        }
                    }

                    total.fetch_add(num_this_tick, SeqCst);
                }

                n
            }));
        }

        let per_thread: Vec<_> = threads.into_iter()
            .map(|th| th.join().unwrap())
            .collect();

        for entry in entries.iter() {
            assert_eq!(PER_ENTRY, entry.num.load(Relaxed));
        }

        for th in per_thread {
            // Kind of annoying that we can't really test anything better than this,
            // but CI tends to be very non deterministic when it comes to multi
            // threading.
            assert!(th > 0, "actual={:?}", th);
        }
    }

    #[test]
    fn with_small_events_collection() {
        const N: usize = 8;
        const ITER: usize = 1_000;

        use std::sync::{Arc, Barrier};
        use std::sync::atomic::AtomicBool;
        use std::sync::atomic::Ordering::{Acquire, Release};
        use std::thread;

        let poll = Poll::new().unwrap();
        let mut registrations = vec![];

        let barrier = Arc::new(Barrier::new(N + 1));
        let done = Arc::new(AtomicBool::new(false));

        for i in 0..N {
            let (registration, set_readiness) = Registration::new2();
            poll.register(&registration, Token(i), Ready::readable(), PollOpt::edge()).unwrap();

            registrations.push(registration);

            let barrier = barrier.clone();
            let done = done.clone();

            thread::spawn(move || {
                barrier.wait();

                while !done.load(Acquire) {
                    set_readiness.set_readiness(Ready::readable()).unwrap();
                }

                // Set one last time
                set_readiness.set_readiness(Ready::readable()).unwrap();
            });
        }

        let mut events = Events::with_capacity(4);

        barrier.wait();

        for _ in 0..ITER {
            poll.poll(&mut events, None).unwrap();
        }

        done.store(true, Release);

        let mut final_ready = vec![false; N];

        for i in 0..5 {
            poll.poll(&mut events, None).unwrap();

            for event in &events {
                final_ready[event.token().0] = true;
            }

            if final_ready.iter().all(|v| *v) {
                return;
            }

            thread::sleep(Duration::from_millis(10));
        }

        panic!("dead lock?");
    }
}

#[test]
fn drop_registration_from_non_main_thread() {
    use std::thread;
    use std::sync::mpsc::channel;

    const THREADS: usize = 8;
    const ITERS: usize = 50_000;

    let mut poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(1024);
    let mut senders = Vec::with_capacity(THREADS);
    let mut token_index = 0;

    // spawn threads, which will send messages to single receiver
    for _ in 0..THREADS {
        let (tx, rx) = channel::<(Registration, SetReadiness)>();
        senders.push(tx);

        thread::spawn(move || {
            for (registration, set_readiness) in rx {
                let _ = set_readiness.set_readiness(Ready::readable());
                drop(registration);
                drop(set_readiness);
            }
        });
    }

    let mut index: usize = 0;
    for _ in 0..ITERS {
        let (registration, set_readiness) = Registration::new(&mut poll, Token(token_index), Ready::readable(), PollOpt::edge());
        let _ = senders[index].send((registration, set_readiness));

        token_index += 1;
        index += 1;
        if index == THREADS {
            index = 0;

            let (registration, set_readiness) = Registration::new(&mut poll, Token(token_index), Ready::readable(), PollOpt::edge());
            let _ = set_readiness.set_readiness(Ready::readable());
            drop(registration);
            drop(set_readiness);
            token_index += 1;

            thread::park_timeout(Duration::from_millis(0));
            let _ = poll.poll(&mut events, None).unwrap();
        }
    }
}
<commit_msg>Remove unused variable warning from tests (#649)<commit_after>use mio::*;
use std::time::Duration;

#[test]
fn smoke() {
    let poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(128);

    let (_r, set) = Registration::new(&poll, Token(0), Ready::readable(), PollOpt::edge());

    let n = poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();
    assert_eq!(n, 0);

    set.set_readiness(Ready::readable()).unwrap();

    let n = poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();
    assert_eq!(n, 1);

    assert_eq!(events.get(0).unwrap().token(), Token(0));
}

#[test]
fn set_readiness_before_register() {
    use std::sync::{Arc, Barrier};
    use std::thread;

    let poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(128);

    for _ in 0..5_000 {
        let (r, set) = Registration::new2();

        let b1 = Arc::new(Barrier::new(2));
        let b2 = b1.clone();

        let th = thread::spawn(move || {
            b2.wait();
            set.set_readiness(Ready::readable()).unwrap();
        });

        b1.wait();

        poll.register(&r, Token(123), Ready::readable(), PollOpt::edge()).unwrap();

        loop {
            let n = poll.poll(&mut events, None).unwrap();

            if n == 0 {
                continue;
            }

            assert_eq!(n, 1);
            assert_eq!(events.get(0).unwrap().token(), Token(123));
            break;
        }

        th.join().unwrap();
    }
}

#[cfg(any(target_os = "linux", target_os = "macos", target_os = "windows"))]
mod stress {
    use mio::*;
    use std::time::Duration;

    #[test]
    fn single_threaded_poll() {
        use std::sync::Arc;
        use std::sync::atomic::AtomicUsize;
        use std::sync::atomic::Ordering::{Acquire, Release};
        use std::thread;

        const NUM_ATTEMPTS: usize = 30;
        const NUM_ITERS: usize = 500;
        const NUM_THREADS: usize = 4;
        const NUM_REGISTRATIONS: usize = 128;

        for _ in 0..NUM_ATTEMPTS {
            let poll = Poll::new().unwrap();
            let mut events = Events::with_capacity(128);

            let registrations: Vec<_> = (0..NUM_REGISTRATIONS).map(|i| {
                Registration::new(&poll, Token(i), Ready::readable(), PollOpt::edge())
            }).collect();

            let mut ready: Vec<_> = (0..NUM_REGISTRATIONS).map(|_| Ready::none()).collect();

            let remaining = Arc::new(AtomicUsize::new(NUM_THREADS));

            for _ in 0..NUM_THREADS {
                let remaining = remaining.clone();

                let set_readiness: Vec<SetReadiness> =
                    registrations.iter().map(|r| r.1.clone()).collect();

                thread::spawn(move || {
                    for _ in 0..NUM_ITERS {
                        for i in 0..NUM_REGISTRATIONS {
                            set_readiness[i].set_readiness(Ready::readable()).unwrap();
                            set_readiness[i].set_readiness(Ready::none()).unwrap();
                            set_readiness[i].set_readiness(Ready::writable()).unwrap();
                            set_readiness[i].set_readiness(Ready::readable() | Ready::writable()).unwrap();
                            set_readiness[i].set_readiness(Ready::none()).unwrap();
                        }
                    }

                    for i in 0..NUM_REGISTRATIONS {
                        set_readiness[i].set_readiness(Ready::readable()).unwrap();
                    }

                    remaining.fetch_sub(1, Release);
                });
            }

            while remaining.load(Acquire) > 0 {
                // Set interest
                for (i, &(ref r, _)) in registrations.iter().enumerate() {
                    r.update(&poll, Token(i), Ready::writable(), PollOpt::edge()).unwrap();
                }

                poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();

                for event in &events {
                    ready[event.token().0] = event.kind();
                }

                // Update registration
                // Set interest
                for (i, &(ref r, _)) in registrations.iter().enumerate() {
                    r.update(&poll, Token(i), Ready::readable(), PollOpt::edge()).unwrap();
                }
            }

            // One final poll
            poll.poll(&mut events, Some(Duration::from_millis(0))).unwrap();

            for event in &events {
                ready[event.token().0] = event.kind();
            }

            // Everything should be flagged as readable
            for ready in ready {
                assert_eq!(ready, Ready::readable());
            }
        }
    }

    #[test]
    fn multi_threaded_poll() {
        use std::sync::{Arc, Barrier};
        use std::sync::atomic::{AtomicUsize};
        use std::sync::atomic::Ordering::{Relaxed, SeqCst};
        use std::thread;

        const ENTRIES: usize = 10_000;
        const PER_ENTRY: usize = 16;
        const THREADS: usize = 4;
        const NUM: usize = ENTRIES * PER_ENTRY;

        struct Entry {
            #[allow(dead_code)]
            registration: Registration,
            set_readiness: SetReadiness,
            num: AtomicUsize,
        }

        impl Entry {
            fn fire(&self) {
                self.set_readiness.set_readiness(Ready::readable()).unwrap();
            }
        }

        let poll = Arc::new(Poll::new().unwrap());
        let mut entries = vec![];

        // Create entries
        for i in 0..ENTRIES {
            let (registration, set_readiness) =
                Registration::new(&poll, Token(i), Ready::readable(), PollOpt::edge());

            entries.push(Entry {
                registration: registration,
                set_readiness: set_readiness,
                num: AtomicUsize::new(0),
            });
        }

        let total = Arc::new(AtomicUsize::new(0));
        let entries = Arc::new(entries);
        let barrier = Arc::new(Barrier::new(THREADS));

        let mut threads = vec![];

        for th in 0..THREADS {
            let poll = poll.clone();
            let total = total.clone();
            let entries = entries.clone();
            let barrier = barrier.clone();

            threads.push(thread::spawn(move || {
                let mut events = Events::with_capacity(128);

                barrier.wait();

                // Prime all the registrations
                let mut i = th;
                while i < ENTRIES {
                    entries[i].fire();
                    i += THREADS;
                }

                let mut n = 0;

                while total.load(SeqCst) < NUM {
                    // A poll timeout is necessary here because there may be more
                    // than one threads blocked in `poll` when the final wakeup
                    // notification arrives (and only notifies one thread).
                    n += poll.poll(&mut events, Some(Duration::from_millis(100))).unwrap();

                    let mut num_this_tick = 0;

                    for event in &events {
                        let e = &entries[event.token().0];

                        let mut num = e.num.load(Relaxed);

                        loop {
                            if num < PER_ENTRY {
                                let actual = e.num.compare_and_swap(num, num + 1, Relaxed);

                                if actual == num {
                                    num_this_tick += 1;
                                    e.fire();
                                    break;
                                }

                                num = actual;
                            } else {
                                break;
                            }
                        }
                    }

                    total.fetch_add(num_this_tick, SeqCst);
                }

                n
            }));
        }

        let per_thread: Vec<_> = threads.into_iter()
            .map(|th| th.join().unwrap())
            .collect();

        for entry in entries.iter() {
            assert_eq!(PER_ENTRY, entry.num.load(Relaxed));
        }

        for th in per_thread {
            // Kind of annoying that we can't really test anything better than this,
            // but CI tends to be very non deterministic when it comes to multi
            // threading.
            assert!(th > 0, "actual={:?}", th);
        }
    }

    #[test]
    fn with_small_events_collection() {
        const N: usize = 8;
        const ITER: usize = 1_000;

        use std::sync::{Arc, Barrier};
        use std::sync::atomic::AtomicBool;
        use std::sync::atomic::Ordering::{Acquire, Release};
        use std::thread;

        let poll = Poll::new().unwrap();
        let mut registrations = vec![];

        let barrier = Arc::new(Barrier::new(N + 1));
        let done = Arc::new(AtomicBool::new(false));

        for i in 0..N {
            let (registration, set_readiness) = Registration::new2();
            poll.register(&registration, Token(i), Ready::readable(), PollOpt::edge()).unwrap();

            registrations.push(registration);

            let barrier = barrier.clone();
            let done = done.clone();

            thread::spawn(move || {
                barrier.wait();

                while !done.load(Acquire) {
                    set_readiness.set_readiness(Ready::readable()).unwrap();
                }

                // Set one last time
                set_readiness.set_readiness(Ready::readable()).unwrap();
            });
        }

        let mut events = Events::with_capacity(4);

        barrier.wait();

        for _ in 0..ITER {
            poll.poll(&mut events, None).unwrap();
        }

        done.store(true, Release);

        let mut final_ready = vec![false; N];

        for _ in 0..5 {
            poll.poll(&mut events, None).unwrap();

            for event in &events {
                final_ready[event.token().0] = true;
            }

            if final_ready.iter().all(|v| *v) {
                return;
            }

            thread::sleep(Duration::from_millis(10));
        }

        panic!("dead lock?");
    }
}

#[test]
fn drop_registration_from_non_main_thread() {
    use std::thread;
    use std::sync::mpsc::channel;

    const THREADS: usize = 8;
    const ITERS: usize = 50_000;

    let mut poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(1024);
    let mut senders = Vec::with_capacity(THREADS);
    let mut token_index = 0;

    // spawn threads, which will send messages to single receiver
    for _ in 0..THREADS {
        let (tx, rx) = channel::<(Registration, SetReadiness)>();
        senders.push(tx);

        thread::spawn(move || {
            for (registration, set_readiness) in rx {
                let _ = set_readiness.set_readiness(Ready::readable());
                drop(registration);
                drop(set_readiness);
            }
        });
    }

    let mut index: usize = 0;
    for _ in 0..ITERS {
        let (registration, set_readiness) = Registration::new(&mut poll, Token(token_index), Ready::readable(), PollOpt::edge());
        let _ = senders[index].send((registration, set_readiness));

        token_index += 1;
        index += 1;
        if index == THREADS {
            index = 0;

            let (registration, set_readiness) = Registration::new(&mut poll, Token(token_index), Ready::readable(), PollOpt::edge());
            let _ = set_readiness.set_readiness(Ready::readable());
            drop(registration);
            drop(set_readiness);
            token_index += 1;

            thread::park_timeout(Duration::from_millis(0));
            let _ = poll.poll(&mut events, None).unwrap();
        }
    }
}
<|endoftext|>


<commit_before>extern crate exonum;
extern crate env_logger;
extern crate clap;
extern crate blockchain_explorer;

use clap::App;

use blockchain_explorer::helpers::GenerateCommand;

fn main() {
    exonum::crypto::init();
    blockchain_explorer::helpers::init_logger().unwrap();

    let app = App::new("Blockchain control utility")
        .version(env!("CARGO_PKG_VERSION"))
        .author("Aleksey S. <aleksei.sidorov@xdev.re>")
        .about("Blockchain control utility")
        .subcommand(GenerateCommand::new());

    let matches = app.get_matches();
    match matches.subcommand() {
        ("generate", Some(matches)) => GenerateCommand::execute(matches),
        _ => unreachable!("Wrong subcommand"),
    }
}<commit_msg>fix after migration<commit_after>extern crate exonum;
extern crate env_logger;
extern crate clap;
extern crate blockchain_explorer;

use clap::App;

use blockchain_explorer::helpers::GenerateCommand;

fn main() {
    exonum::crypto::init();
    blockchain_explorer::helpers::init_logger().unwrap();

    let app = App::new("Blockchain control utility")
        .version(env!("CARGO_PKG_VERSION"))
        .author("Aleksey S. <aleksei.sidorov@xdevre>")
        .about("Blockchain control utility")
        .subcommand(GenerateCommand::new());

    let matches = app.get_matches();
    match matches.subcommand() {
        ("generate", Some(matches)) => GenerateCommand::execute(matches),
        _ => unreachable!("Wrong subcommand"),
    }
}<|endoftext|>


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

use std::mem;
use std::vec::Vec;

extern "C" {
    fn realloc(ptr: *mut u8, bytes: usize) -> *mut u8;
    fn malloc(bytes: usize) -> *mut u8;
}

pub trait FallibleVec<T> {
    /// Append |val| to the end of |vec|.  Returns Ok(()) on success,
    /// Err(()) if it fails, which can only be due to lack of memory.
    fn try_push(&mut self, value: T) -> Result<(), ()>;

    /// Expand the vector size. Return Ok(()) on success, Err(()) if it
    /// fails.
    fn try_reserve(&mut self, new_cap: usize) -> Result<(), ()>;
}

/////////////////////////////////////////////////////////////////
// Vec

impl<T> FallibleVec<T> for Vec<T> {
    #[inline]
    fn try_push(&mut self, val: T) -> Result<(), ()> {
        if self.capacity() == self.len() {
            let old_cap: usize = self.capacity();
            let new_cap: usize
                = if old_cap == 0 { 4 } else { old_cap.checked_mul(2).ok_or(()) ? };

            try_extend_vec(self, new_cap)?;
            debug_assert!(self.capacity() > self.len());
        }
        self.push(val);
        Ok(())
    }

    #[inline]
    fn try_reserve(&mut self, cap: usize) -> Result<(), ()> {
        let new_cap = cap + self.capacity();
        try_extend_vec(self, new_cap)?;
        debug_assert!(self.capacity() == new_cap);
        Ok(())
    }
}

#[inline(never)]
#[cold]
fn try_extend_vec<T>(vec: &mut Vec<T>, new_cap: usize) -> Result<(), ()> {
    let old_ptr = vec.as_mut_ptr();
    let old_len = vec.len();

    let old_cap: usize = vec.capacity();

    if old_cap > new_cap {
        return Ok(());
    }

    let new_size_bytes
        = new_cap.checked_mul(mem::size_of::<T>()).ok_or(()) ? ;

    let new_ptr = unsafe {
        if old_cap == 0 {
            malloc(new_size_bytes)
        } else {
            realloc(old_ptr as *mut u8, new_size_bytes)
        }
    };

    if new_ptr.is_null() {
        return Err(());
    }

    let new_vec = unsafe {
        Vec::from_raw_parts(new_ptr as *mut T, old_len, new_cap)
    };

    mem::forget(mem::replace(vec, new_vec));
    Ok(())
}

#[test]
fn oom_test() {
    let mut vec: Vec<char> = Vec::new();
    match vec.try_reserve(std::usize::MAX) {
        Ok(_) => panic!("it should be OOM"),
        _ => (),
    }
}
<commit_msg>Don't (re)allocate if the new_cap is the same size as old_cap.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::mem;
use std::vec::Vec;

extern "C" {
    fn realloc(ptr: *mut u8, bytes: usize) -> *mut u8;
    fn malloc(bytes: usize) -> *mut u8;
}

pub trait FallibleVec<T> {
    /// Append |val| to the end of |vec|.  Returns Ok(()) on success,
    /// Err(()) if it fails, which can only be due to lack of memory.
    fn try_push(&mut self, value: T) -> Result<(), ()>;

    /// Expand the vector size. Return Ok(()) on success, Err(()) if it
    /// fails.
    fn try_reserve(&mut self, new_cap: usize) -> Result<(), ()>;
}

/////////////////////////////////////////////////////////////////
// Vec

impl<T> FallibleVec<T> for Vec<T> {
    #[inline]
    fn try_push(&mut self, val: T) -> Result<(), ()> {
        if self.capacity() == self.len() {
            let old_cap: usize = self.capacity();
            let new_cap: usize
                = if old_cap == 0 { 4 } else { old_cap.checked_mul(2).ok_or(()) ? };

            try_extend_vec(self, new_cap)?;
            debug_assert!(self.capacity() > self.len());
        }
        self.push(val);
        Ok(())
    }

    #[inline]
    fn try_reserve(&mut self, cap: usize) -> Result<(), ()> {
        let new_cap = cap + self.capacity();
        try_extend_vec(self, new_cap)?;
        debug_assert!(self.capacity() == new_cap);
        Ok(())
    }
}

#[inline(never)]
#[cold]
fn try_extend_vec<T>(vec: &mut Vec<T>, new_cap: usize) -> Result<(), ()> {
    let old_ptr = vec.as_mut_ptr();
    let old_len = vec.len();

    let old_cap: usize = vec.capacity();

    if old_cap >= new_cap {
        return Ok(());
    }

    let new_size_bytes
        = new_cap.checked_mul(mem::size_of::<T>()).ok_or(()) ? ;

    let new_ptr = unsafe {
        if old_cap == 0 {
            malloc(new_size_bytes)
        } else {
            realloc(old_ptr as *mut u8, new_size_bytes)
        }
    };

    if new_ptr.is_null() {
        return Err(());
    }

    let new_vec = unsafe {
        Vec::from_raw_parts(new_ptr as *mut T, old_len, new_cap)
    };

    mem::forget(mem::replace(vec, new_vec));
    Ok(())
}

#[test]
fn oom_test() {
    let mut vec: Vec<char> = Vec::new();
    match vec.try_reserve(std::usize::MAX) {
        Ok(_) => panic!("it should be OOM"),
        _ => (),
    }
}
<|endoftext|>


<commit_before>use std::borrow::Cow;
use std::env;
use std::ffi::OsStr;
use std::ffi::OsString;
use std::path::{Path, PathBuf};
use std::process::Command;

macro_rules! forward_features {
	($cmd:expr, $($feature:literal,)+ ) => {
		let mut features = vec![];

		$(
			if cfg!(feature = $feature) {
				features.push($feature);
			}
		)+

		if !features.is_empty() {
			$cmd.arg("--features");
			$cmd.arg(features.join(" "));
		}
	};
}

fn build_hermit(src_dir: &Path) {
	let target_dir = {
		let mut target_dir = PathBuf::from(env::var_os("OUT_DIR").unwrap());
		target_dir.push("target");
		target_dir
	};
	let manifest_path = src_dir.join("Cargo.toml");
	assert!(
		manifest_path.exists(),
		"kernel manifest path `{}` does not exist",
		manifest_path.display()
	);
	let target_arch = env::var("CARGO_CFG_TARGET_ARCH").unwrap();
	let profile = env::var("PROFILE").expect("PROFILE was not set");
	let mut cmd = Command::new("cargo");

	let kernel_triple = match target_arch.as_str() {
		"x86_64" => "x86_64-unknown-none-hermitkernel",
		"aarch64" => "aarch64-unknown-hermit",
		_ => panic!("Unsupported target arch: {}", target_arch),
	};

	cmd.arg("build")
		.arg("-Z")
		.arg("build-std=core,alloc")
		.arg("--target")
		.arg(kernel_triple)
		.arg("--manifest-path")
		.arg("Cargo.toml");

	cmd.current_dir(src_dir);

	cmd.env_remove("RUSTUP_TOOLCHAIN");
	if option_env!("RUSTC_WORKSPACE_WRAPPER")
		.unwrap_or_default()
		.ends_with("clippy-driver")
	{
		cmd.env("RUSTC_WORKSPACE_WRAPPER", "clippy-driver");
	}

	cmd.env("CARGO_TERM_COLOR", "always");

	cmd.arg("--target-dir").arg(&target_dir);

	if profile == "release" {
		cmd.arg("--release");
	}

	// Control enabled features via this crate's features
	cmd.arg("--no-default-features");
	forward_features!(cmd, "acpi", "fsgsbase", "pci", "smp", "vga",);

	let mut rustflags = vec!["-Zmutable-noalias=no".to_string()];
	let outer_rustflags = env::var("CARGO_ENCODED_RUSTFLAGS").unwrap();

	#[cfg(feature = "instrument")]
	{
		rustflags.push("-Zinstrument-mcount".to_string());
		// Add outer rustflags to command
		rustflags.push(outer_rustflags);
	}

	#[cfg(not(feature = "instrument"))]
	{
		// If the `instrument` feature feature is not enabled,
		// filter it from outer rustflags before adding them to the command.
		if !outer_rustflags.is_empty() {
			let flags = outer_rustflags
				.split('\x1f')
				.filter(|&flag| !flag.contains("instrument-mcount"))
				.map(String::from);
			rustflags.extend(flags);
		}
	}

	cmd.env("CARGO_ENCODED_RUSTFLAGS", rustflags.join("\x1f"));

	let status = cmd.status().expect("failed to start kernel build");
	assert!(status.success());

	let lib_location = target_dir
		.join(kernel_triple)
		.join(&profile)
		.canonicalize()
		.unwrap();

	println!("Lib location: {}", lib_location.display());

	let lib = lib_location.join("libhermit.a");

	let mut symbols = vec!["rust_begin_unwind", "rust_oom"];

	if target_arch == "aarch64" {
		symbols.extend(include_str!("aarch64-duplicate-symbols").lines());
	}

	rename_symbols(symbols, &lib);

	println!("cargo:rustc-link-search=native={}", lib_location.display());
	println!("cargo:rustc-link-lib=static=hermit");

	println!("cargo:rerun-if-changed={}", src_dir.display());
	// HERMIT_LOG_LEVEL_FILTER sets the log level filter at compile time
	println!("cargo:rerun-if-env-changed=HERMIT_LOG_LEVEL_FILTER");
}

/// Kernel and user space has its own versions of panic handler, oom handler, memcpy, memset, etc,
/// Consequently, we rename the functions in the libos to avoid collisions.
fn rename_symbols(symbols: impl IntoIterator<Item = impl AsRef<OsStr>>, lib: impl AsRef<Path>) {
	let args = symbols.into_iter().flat_map(|symbol| {
		let option = OsStr::new("--redefine-sym");
		let arg = [symbol.as_ref(), "=kernel-".as_ref(), symbol.as_ref()]
			.into_iter()
			.collect::<OsString>();
		[Cow::Borrowed(option), Cow::Owned(arg)]
	});

	let status = Command::new("rust-objcopy")
		.args(args)
		.arg(lib.as_ref())
		.status()
		.expect("Failed to execute rust-objcopy. Is cargo-binutils installed?");
	assert!(status.success(), "rust-objcopy was not successful");
}

#[cfg(not(feature = "with_submodule"))]
fn build() {
	#[cfg(windows)]
	let out_dir = env::temp_dir();
	#[cfg(not(windows))]
	let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
	let src_dir = out_dir.join("rusty-hermit");

	if !src_dir.as_path().exists() {
		let status = Command::new(env!("CARGO"))
			.current_dir(out_dir)
			.arg("download")
			.arg("--output")
			.arg(src_dir.clone().into_os_string())
			.arg("--extract")
			.arg("rusty-hermit")
			.status()
			.expect("failed to start kernel download");
		assert!(
			status.success(),
			"Unable to download rusty-hermit. Is cargo-download installed?"
		);
	}

	build_hermit(src_dir.as_ref());
}

#[cfg(feature = "with_submodule")]
fn build() {
	let src_dir = env::current_dir()
		.unwrap()
		.parent()
		.unwrap()
		.join("libhermit-rs");

	build_hermit(src_dir.as_ref());
}

fn main() {
	build();
}
<commit_msg>build.rs: Don't special-case Windows download dir<commit_after>use std::borrow::Cow;
use std::env;
use std::ffi::OsStr;
use std::ffi::OsString;
use std::path::{Path, PathBuf};
use std::process::Command;

macro_rules! forward_features {
	($cmd:expr, $($feature:literal,)+ ) => {
		let mut features = vec![];

		$(
			if cfg!(feature = $feature) {
				features.push($feature);
			}
		)+

		if !features.is_empty() {
			$cmd.arg("--features");
			$cmd.arg(features.join(" "));
		}
	};
}

fn build_hermit(src_dir: &Path) {
	let target_dir = {
		let mut target_dir = PathBuf::from(env::var_os("OUT_DIR").unwrap());
		target_dir.push("target");
		target_dir
	};
	let manifest_path = src_dir.join("Cargo.toml");
	assert!(
		manifest_path.exists(),
		"kernel manifest path `{}` does not exist",
		manifest_path.display()
	);
	let target_arch = env::var("CARGO_CFG_TARGET_ARCH").unwrap();
	let profile = env::var("PROFILE").expect("PROFILE was not set");
	let mut cmd = Command::new("cargo");

	let kernel_triple = match target_arch.as_str() {
		"x86_64" => "x86_64-unknown-none-hermitkernel",
		"aarch64" => "aarch64-unknown-hermit",
		_ => panic!("Unsupported target arch: {}", target_arch),
	};

	cmd.arg("build")
		.arg("-Z")
		.arg("build-std=core,alloc")
		.arg("--target")
		.arg(kernel_triple)
		.arg("--manifest-path")
		.arg("Cargo.toml");

	cmd.current_dir(src_dir);

	cmd.env_remove("RUSTUP_TOOLCHAIN");
	if option_env!("RUSTC_WORKSPACE_WRAPPER")
		.unwrap_or_default()
		.ends_with("clippy-driver")
	{
		cmd.env("RUSTC_WORKSPACE_WRAPPER", "clippy-driver");
	}

	cmd.env("CARGO_TERM_COLOR", "always");

	cmd.arg("--target-dir").arg(&target_dir);

	if profile == "release" {
		cmd.arg("--release");
	}

	// Control enabled features via this crate's features
	cmd.arg("--no-default-features");
	forward_features!(cmd, "acpi", "fsgsbase", "pci", "smp", "vga",);

	let mut rustflags = vec!["-Zmutable-noalias=no".to_string()];
	let outer_rustflags = env::var("CARGO_ENCODED_RUSTFLAGS").unwrap();

	#[cfg(feature = "instrument")]
	{
		rustflags.push("-Zinstrument-mcount".to_string());
		// Add outer rustflags to command
		rustflags.push(outer_rustflags);
	}

	#[cfg(not(feature = "instrument"))]
	{
		// If the `instrument` feature feature is not enabled,
		// filter it from outer rustflags before adding them to the command.
		if !outer_rustflags.is_empty() {
			let flags = outer_rustflags
				.split('\x1f')
				.filter(|&flag| !flag.contains("instrument-mcount"))
				.map(String::from);
			rustflags.extend(flags);
		}
	}

	cmd.env("CARGO_ENCODED_RUSTFLAGS", rustflags.join("\x1f"));

	let status = cmd.status().expect("failed to start kernel build");
	assert!(status.success());

	let lib_location = target_dir
		.join(kernel_triple)
		.join(&profile)
		.canonicalize()
		.unwrap();

	println!("Lib location: {}", lib_location.display());

	let lib = lib_location.join("libhermit.a");

	let mut symbols = vec!["rust_begin_unwind", "rust_oom"];

	if target_arch == "aarch64" {
		symbols.extend(include_str!("aarch64-duplicate-symbols").lines());
	}

	rename_symbols(symbols, &lib);

	println!("cargo:rustc-link-search=native={}", lib_location.display());
	println!("cargo:rustc-link-lib=static=hermit");

	println!("cargo:rerun-if-changed={}", src_dir.display());
	// HERMIT_LOG_LEVEL_FILTER sets the log level filter at compile time
	println!("cargo:rerun-if-env-changed=HERMIT_LOG_LEVEL_FILTER");
}

/// Kernel and user space has its own versions of panic handler, oom handler, memcpy, memset, etc,
/// Consequently, we rename the functions in the libos to avoid collisions.
fn rename_symbols(symbols: impl IntoIterator<Item = impl AsRef<OsStr>>, lib: impl AsRef<Path>) {
	let args = symbols.into_iter().flat_map(|symbol| {
		let option = OsStr::new("--redefine-sym");
		let arg = [symbol.as_ref(), "=kernel-".as_ref(), symbol.as_ref()]
			.into_iter()
			.collect::<OsString>();
		[Cow::Borrowed(option), Cow::Owned(arg)]
	});

	let status = Command::new("rust-objcopy")
		.args(args)
		.arg(lib.as_ref())
		.status()
		.expect("Failed to execute rust-objcopy. Is cargo-binutils installed?");
	assert!(status.success(), "rust-objcopy was not successful");
}

#[cfg(not(feature = "with_submodule"))]
fn build() {
	let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
	let src_dir = out_dir.join("rusty-hermit");

	if !src_dir.as_path().exists() {
		let status = Command::new(env!("CARGO"))
			.current_dir(out_dir)
			.arg("download")
			.arg("--output")
			.arg(src_dir.clone().into_os_string())
			.arg("--extract")
			.arg("rusty-hermit")
			.status()
			.expect("failed to start kernel download");
		assert!(
			status.success(),
			"Unable to download rusty-hermit. Is cargo-download installed?"
		);
	}

	build_hermit(src_dir.as_ref());
}

#[cfg(feature = "with_submodule")]
fn build() {
	let src_dir = env::current_dir()
		.unwrap()
		.parent()
		.unwrap()
		.join("libhermit-rs");

	build_hermit(src_dir.as_ref());
}

fn main() {
	build();
}
<|endoftext|>


<commit_before>//! Aim to provide platform-independent methods to obtain login records
//!
//! **ONLY** support linux, macos and freebsd for the time being

// This file is part of the uutils coreutils package.
//
// (c) Jian Zeng <anonymousknight96@gmail.com>
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.
//

use super::libc;
pub extern crate time;
use self::time::{Tm, Timespec};

use ::std::io::Result as IOResult;
use ::std::io::Error as IOError;
use ::std::ptr;
use ::std::ffi::CString;

pub use self::ut::*;
use libc::utmpx;
// pub use libc::getutxid;
// pub use libc::getutxline;
// pub use libc::pututxline;
pub use libc::getutxent;
pub use libc::setutxent;
pub use libc::endutxent;
#[cfg(any(target_os = "macos", target_os = "linux"))]
pub use libc::utmpxname;
#[cfg(target_os = "freebsd")]
pub unsafe extern "C" fn utmpxname(_file: *const libc::c_char) -> libc::c_int {
    0
}

// In case the c_char array doesn' t end with NULL
macro_rules! chars2string {
    ($arr:expr) => (
        $arr.iter().take_while(|i| **i > 0).map(|&i| i as u8 as char).collect::<String>()
    )
}

#[cfg(target_os = "linux")]
mod ut {
    pub static DEFAULT_FILE: &'static str = "/var/run/utmp";

    pub use libc::__UT_LINESIZE as UT_LINESIZE;
    pub use libc::__UT_NAMESIZE as UT_NAMESIZE;
    pub use libc::__UT_HOSTSIZE as UT_HOSTSIZE;
    pub const UT_IDSIZE: usize = 4;

    pub use libc::EMPTY;
    pub use libc::RUN_LVL;
    pub use libc::BOOT_TIME;
    pub use libc::NEW_TIME;
    pub use libc::OLD_TIME;
    pub use libc::INIT_PROCESS;
    pub use libc::LOGIN_PROCESS;
    pub use libc::USER_PROCESS;
    pub use libc::DEAD_PROCESS;
    pub use libc::ACCOUNTING;
}

#[cfg(target_os = "macos")]
mod ut {
    pub static DEFAULT_FILE: &'static str = "/var/run/utmpx";

    pub use libc::_UTX_LINESIZE as UT_LINESIZE;
    pub use libc::_UTX_USERSIZE as UT_NAMESIZE;
    pub use libc::_UTX_HOSTSIZE as UT_HOSTSIZE;
    pub use libc::_UTX_IDSIZE as UT_IDSIZE;

    pub use libc::EMPTY;
    pub use libc::RUN_LVL;
    pub use libc::BOOT_TIME;
    pub use libc::NEW_TIME;
    pub use libc::OLD_TIME;
    pub use libc::INIT_PROCESS;
    pub use libc::LOGIN_PROCESS;
    pub use libc::USER_PROCESS;
    pub use libc::DEAD_PROCESS;
    pub use libc::ACCOUNTING;
    pub use libc::SIGNATURE;
    pub use libc::SHUTDOWN_TIME;
}

#[cfg(target_os = "freebsd")]
mod ut {
    use super::libc;

    pub static DEFAULT_FILE: &'static str = "";

    pub const UT_LINESIZE: usize = 16;
    pub const UT_NAMESIZE: usize = 32;
    pub const UT_IDSIZE: usize = 8;
    pub const UT_HOSTSIZE: usize = 128;

    pub use libc::EMPTY;
    pub use libc::BOOT_TIME;
    pub use libc::OLD_TIME;
    pub use libc::NEW_TIME;
    pub use libc::USER_PROCESS;
    pub use libc::INIT_PROCESS;
    pub use libc::LOGIN_PROCESS;
    pub use libc::DEAD_PROCESS;
    pub use libc::SHUTDOWN_TIME;
}

/// Login records
///
/// Examples:
/// ---------
///
/// ```
/// for ut in Utmpx::iter_all_records() {
///     if ut.is_user_process() {
///         println!("{}: {}", ut.host(), ut.user())
///     }
/// }
/// ```
///
/// Specifying the path to login record:
///
/// ```
/// for ut in Utmpx::iter_all_records().read_from("/some/where/else") {
///     if ut.is_user_process() {
///         println!("{}: {}", ut.host(), ut.user())
///     }
/// }
/// ```
pub struct Utmpx {
    inner: utmpx,
}

impl Utmpx {
    /// A.K.A. ut.ut_type
    pub fn record_type(&self) -> i16 {
        self.inner.ut_type as i16
    }
    /// A.K.A. ut.ut_pid
    pub fn pid(&self) -> i32 {
        self.inner.ut_pid as i32
    }
    /// A.K.A. ut.ut_id
    pub fn terminal_suffix(&self) -> String {
        chars2string!(self.inner.ut_id)
    }
    /// A.K.A. ut.ut_user
    pub fn user(&self) -> String {
        chars2string!(self.inner.ut_user)
    }
    /// A.K.A. ut.ut_host
    pub fn host(&self) -> String {
        chars2string!(self.inner.ut_host)
    }
    /// A.K.A. ut.ut_line
    pub fn tty_device(&self) -> String {
        chars2string!(self.inner.ut_line)
    }
    /// A.K.A. ut.ut_tv
    pub fn login_time(&self) -> Tm {
        time::at(Timespec::new(self.inner.ut_tv.tv_sec as i64,
                               self.inner.ut_tv.tv_usec as i32))
    }
    /// A.K.A. ut.ut_exit
    ///
    /// Return (e_termination, e_exit)
    #[cfg(any(target_os = "linux", target_os = "android"))]
    pub fn exit_status(&self) -> (i16, i16) {
        (self.inner.ut_exit.e_termination, self.inner.ut_exit.e_exit)
    }
    /// A.K.A. ut.ut_exit
    ///
    /// Return (0, 0) on Non-Linux platform
    #[cfg(not(any(target_os = "linux", target_os = "android")))]
    pub fn exit_status(&self) -> (i16, i16) {
        (0, 0)
    }
    /// Consumes the `Utmpx`, returning the underlying C struct utmpx
    pub fn into_inner(self) -> utmpx {
        self.inner
    }
    pub fn is_user_process(&self) -> bool {
        !self.user().is_empty() && self.record_type() == USER_PROCESS
    }

    /// Canonicalize host name using DNS
    pub fn canon_host(&self) -> IOResult<String> {
        const AI_CANONNAME: libc::c_int = 0x2;
        let host = self.host();
        let host = host.split(':').nth(0).unwrap();
        let hints = libc::addrinfo {
            ai_flags: AI_CANONNAME,
            ai_family: 0,
            ai_socktype: 0,
            ai_protocol: 0,
            ai_addrlen: 0,
            ai_addr: ptr::null_mut(),
            ai_canonname: ptr::null_mut(),
            ai_next: ptr::null_mut(),
        };
        let c_host = CString::new(host).unwrap();
        let mut res = ptr::null_mut();
        let status = unsafe {
            libc::getaddrinfo(c_host.as_ptr(),
                              ptr::null(),
                              &hints as *const _,
                              &mut res as *mut _)
        };
        if status == 0 {
            let info: libc::addrinfo = unsafe { ptr::read(res as *const _) };
            // http://lists.gnu.org/archive/html/bug-coreutils/2006-09/msg00300.html
            // says Darwin 7.9.0 getaddrinfo returns 0 but sets
            // res->ai_canonname to NULL.
            let ret = if info.ai_canonname.is_null() {
                Ok(String::from(host))
            } else {
                Ok(unsafe { CString::from_raw(info.ai_canonname).into_string().unwrap() })
            };
            unsafe {
                libc::freeaddrinfo(res);
            }
            ret
        } else {
            Err(IOError::last_os_error())
        }
    }
    pub fn iter_all_records() -> UtmpxIter {
        UtmpxIter
    }
}

/// Iterator of login records
pub struct UtmpxIter;

impl UtmpxIter {
    /// Sets the name of the utmpx-format file for the other utmpx functions to access.
    ///
    /// If not set, default record file will be used(file path depends on the target OS)
    pub fn read_from(self, f: &str) -> Self {
        let res = unsafe { utmpxname(CString::new(f).unwrap().as_ptr()) };
        if res != 0 {
            println!("Warning: {}", IOError::last_os_error());
        }
        unsafe {
            setutxent();
        }
        self
    }
}

impl Iterator for UtmpxIter {
    type Item = Utmpx;
    fn next(&mut self) -> Option<Self::Item> {
        unsafe {
            let res = getutxent();
            if !res.is_null() {
                Some(Utmpx { inner: ptr::read(res as *const _) })
            } else {
                endutxent();
                None
            }
        }
    }
}
<commit_msg>uucore::utmpx: fix doctest errors<commit_after>// This file is part of the uutils coreutils package.
//
// (c) Jian Zeng <anonymousknight96@gmail.com>
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.
//
//! Aims to provide platform-independent methods to obtain login records
//!
//! **ONLY** support linux, macos and freebsd for the time being
//!
//! # Examples:
//!
//! ```
//! use uucore::utmpx::Utmpx;
//! for ut in Utmpx::iter_all_records() {
//!     if ut.is_user_process() {
//!         println!("{}: {}", ut.host(), ut.user())
//!     }
//! }
//! ```
//!
//! Specifying the path to login record:
//!
//! ```
//! use uucore::utmpx::Utmpx;
//! for ut in Utmpx::iter_all_records().read_from("/some/where/else") {
//!     if ut.is_user_process() {
//!         println!("{}: {}", ut.host(), ut.user())
//!     }
//! }
//! ```

use super::libc;
pub extern crate time;
use self::time::{Tm, Timespec};

use ::std::io::Result as IOResult;
use ::std::io::Error as IOError;
use ::std::ptr;
use ::std::ffi::CString;

pub use self::ut::*;
use libc::utmpx;
// pub use libc::getutxid;
// pub use libc::getutxline;
// pub use libc::pututxline;
pub use libc::getutxent;
pub use libc::setutxent;
pub use libc::endutxent;
#[cfg(any(target_os = "macos", target_os = "linux"))]
pub use libc::utmpxname;
#[cfg(target_os = "freebsd")]
pub unsafe extern "C" fn utmpxname(_file: *const libc::c_char) -> libc::c_int {
    0
}

// In case the c_char array doesn' t end with NULL
macro_rules! chars2string {
    ($arr:expr) => (
        $arr.iter().take_while(|i| **i > 0).map(|&i| i as u8 as char).collect::<String>()
    )
}

#[cfg(target_os = "linux")]
mod ut {
    pub static DEFAULT_FILE: &'static str = "/var/run/utmp";

    pub use libc::__UT_LINESIZE as UT_LINESIZE;
    pub use libc::__UT_NAMESIZE as UT_NAMESIZE;
    pub use libc::__UT_HOSTSIZE as UT_HOSTSIZE;
    pub const UT_IDSIZE: usize = 4;

    pub use libc::EMPTY;
    pub use libc::RUN_LVL;
    pub use libc::BOOT_TIME;
    pub use libc::NEW_TIME;
    pub use libc::OLD_TIME;
    pub use libc::INIT_PROCESS;
    pub use libc::LOGIN_PROCESS;
    pub use libc::USER_PROCESS;
    pub use libc::DEAD_PROCESS;
    pub use libc::ACCOUNTING;
}

#[cfg(target_os = "macos")]
mod ut {
    pub static DEFAULT_FILE: &'static str = "/var/run/utmpx";

    pub use libc::_UTX_LINESIZE as UT_LINESIZE;
    pub use libc::_UTX_USERSIZE as UT_NAMESIZE;
    pub use libc::_UTX_HOSTSIZE as UT_HOSTSIZE;
    pub use libc::_UTX_IDSIZE as UT_IDSIZE;

    pub use libc::EMPTY;
    pub use libc::RUN_LVL;
    pub use libc::BOOT_TIME;
    pub use libc::NEW_TIME;
    pub use libc::OLD_TIME;
    pub use libc::INIT_PROCESS;
    pub use libc::LOGIN_PROCESS;
    pub use libc::USER_PROCESS;
    pub use libc::DEAD_PROCESS;
    pub use libc::ACCOUNTING;
    pub use libc::SIGNATURE;
    pub use libc::SHUTDOWN_TIME;
}

#[cfg(target_os = "freebsd")]
mod ut {
    use super::libc;

    pub static DEFAULT_FILE: &'static str = "";

    pub const UT_LINESIZE: usize = 16;
    pub const UT_NAMESIZE: usize = 32;
    pub const UT_IDSIZE: usize = 8;
    pub const UT_HOSTSIZE: usize = 128;

    pub use libc::EMPTY;
    pub use libc::BOOT_TIME;
    pub use libc::OLD_TIME;
    pub use libc::NEW_TIME;
    pub use libc::USER_PROCESS;
    pub use libc::INIT_PROCESS;
    pub use libc::LOGIN_PROCESS;
    pub use libc::DEAD_PROCESS;
    pub use libc::SHUTDOWN_TIME;
}

pub struct Utmpx {
    inner: utmpx,
}

impl Utmpx {
    /// A.K.A. ut.ut_type
    pub fn record_type(&self) -> i16 {
        self.inner.ut_type as i16
    }
    /// A.K.A. ut.ut_pid
    pub fn pid(&self) -> i32 {
        self.inner.ut_pid as i32
    }
    /// A.K.A. ut.ut_id
    pub fn terminal_suffix(&self) -> String {
        chars2string!(self.inner.ut_id)
    }
    /// A.K.A. ut.ut_user
    pub fn user(&self) -> String {
        chars2string!(self.inner.ut_user)
    }
    /// A.K.A. ut.ut_host
    pub fn host(&self) -> String {
        chars2string!(self.inner.ut_host)
    }
    /// A.K.A. ut.ut_line
    pub fn tty_device(&self) -> String {
        chars2string!(self.inner.ut_line)
    }
    /// A.K.A. ut.ut_tv
    pub fn login_time(&self) -> Tm {
        time::at(Timespec::new(self.inner.ut_tv.tv_sec as i64,
                               self.inner.ut_tv.tv_usec as i32))
    }
    /// A.K.A. ut.ut_exit
    ///
    /// Return (e_termination, e_exit)
    #[cfg(any(target_os = "linux", target_os = "android"))]
    pub fn exit_status(&self) -> (i16, i16) {
        (self.inner.ut_exit.e_termination, self.inner.ut_exit.e_exit)
    }
    /// A.K.A. ut.ut_exit
    ///
    /// Return (0, 0) on Non-Linux platform
    #[cfg(not(any(target_os = "linux", target_os = "android")))]
    pub fn exit_status(&self) -> (i16, i16) {
        (0, 0)
    }
    /// Consumes the `Utmpx`, returning the underlying C struct utmpx
    pub fn into_inner(self) -> utmpx {
        self.inner
    }
    pub fn is_user_process(&self) -> bool {
        !self.user().is_empty() && self.record_type() == USER_PROCESS
    }

    /// Canonicalize host name using DNS
    pub fn canon_host(&self) -> IOResult<String> {
        const AI_CANONNAME: libc::c_int = 0x2;
        let host = self.host();
        let host = host.split(':').nth(0).unwrap();
        let hints = libc::addrinfo {
            ai_flags: AI_CANONNAME,
            ai_family: 0,
            ai_socktype: 0,
            ai_protocol: 0,
            ai_addrlen: 0,
            ai_addr: ptr::null_mut(),
            ai_canonname: ptr::null_mut(),
            ai_next: ptr::null_mut(),
        };
        let c_host = CString::new(host).unwrap();
        let mut res = ptr::null_mut();
        let status = unsafe {
            libc::getaddrinfo(c_host.as_ptr(),
                              ptr::null(),
                              &hints as *const _,
                              &mut res as *mut _)
        };
        if status == 0 {
            let info: libc::addrinfo = unsafe { ptr::read(res as *const _) };
            // http://lists.gnu.org/archive/html/bug-coreutils/2006-09/msg00300.html
            // says Darwin 7.9.0 getaddrinfo returns 0 but sets
            // res->ai_canonname to NULL.
            let ret = if info.ai_canonname.is_null() {
                Ok(String::from(host))
            } else {
                Ok(unsafe { CString::from_raw(info.ai_canonname).into_string().unwrap() })
            };
            unsafe {
                libc::freeaddrinfo(res);
            }
            ret
        } else {
            Err(IOError::last_os_error())
        }
    }
    pub fn iter_all_records() -> UtmpxIter {
        UtmpxIter
    }
}

/// Iterator of login records
pub struct UtmpxIter;

impl UtmpxIter {
    /// Sets the name of the utmpx-format file for the other utmpx functions to access.
    ///
    /// If not set, default record file will be used(file path depends on the target OS)
    pub fn read_from(self, f: &str) -> Self {
        let res = unsafe { utmpxname(CString::new(f).unwrap().as_ptr()) };
        if res != 0 {
            println!("Warning: {}", IOError::last_os_error());
        }
        unsafe {
            setutxent();
        }
        self
    }
}

impl Iterator for UtmpxIter {
    type Item = Utmpx;
    fn next(&mut self) -> Option<Self::Item> {
        unsafe {
            let res = getutxent();
            if !res.is_null() {
                Some(Utmpx { inner: ptr::read(res as *const _) })
            } else {
                endutxent();
                None
            }
        }
    }
}
<|endoftext|>


<commit_before>use ::*;

pub struct Timer {
    time : Duration,
    time_left : Duration,
    pub ready : bool,
}

impl Timer {
    pub fn from_millis(milliseconds : u64) -> Self {
        let duration = Duration::from_millis(milliseconds);
        Self {
            time : duration,
            time_left : duration,
            ready : false,
        }
    }

    pub fn reset(&mut self) {
        self.ready = false;
        self.time_left = self.time;
    }

    pub fn update(&mut self, delta : Duration) {
        if self.ready {
            return;
        }
        if let Some(result) = self.time_left.checked_sub(delta) {
            self.time_left = result;
        } else {
            self.ready = true;
        }
    }
}
<commit_msg>shorten a variable<commit_after>use ::*;

pub struct Timer {
    time : Duration,
    time_left : Duration,
    pub ready : bool,
}

impl Timer {
    pub fn from_millis(ms : u64) -> Self {
        let duration = Duration::from_millis(ms);
        Self {
            time : duration,
            time_left : duration,
            ready : false,
        }
    }

    pub fn reset(&mut self) {
        self.ready = false;
        self.time_left = self.time;
    }

    pub fn update(&mut self, delta : Duration) {
        if self.ready {
            return;
        }
        if let Some(result) = self.time_left.checked_sub(delta) {
            self.time_left = result;
        } else {
            self.ready = true;
        }
    }
}
<|endoftext|>


<commit_before>use std::sync::mpsc;
use std::thread;

use super::{Event, MiliSec, Request};

#[derive(Debug)]
pub struct TimerEvent(pub MiliSec);

pub struct Timer;

impl Timer {
    pub fn from_chan(ms: MiliSec, tx: mpsc::Sender<Request>) {
        thread::spawn(move || {
            loop {
                thread::sleep_ms(ms);
                if tx.send(Request::Event(Event::Timer(TimerEvent(ms)))).is_err() {
                    break;
                }
            }
        });
    }
}
<commit_msg>timer: mark TimerEvent as Copy<commit_after>use std::sync::mpsc;
use std::thread;

use super::{Event, MiliSec, Request};

#[derive(Clone, Copy, Debug)]
pub struct TimerEvent(pub MiliSec);

pub struct Timer;

impl Timer {
    pub fn from_chan(ms: MiliSec, tx: mpsc::Sender<Request>) {
        thread::spawn(move || {
            loop {
                thread::sleep_ms(ms);
                if tx.send(Request::Event(Event::Timer(TimerEvent(ms)))).is_err() {
                    break;
                }
            }
        });
    }
}
<|endoftext|>


<commit_before>// Take a look at the license at the top of the repository in the LICENSE file.

use crate::{DiskExt, DiskType};

use std::ffi::{OsStr, OsString};
use std::mem::size_of;
use std::path::Path;

use winapi::ctypes::c_void;
use winapi::shared::minwindef::{DWORD, MAX_PATH};
use winapi::um::fileapi::{
    CreateFileW, GetDiskFreeSpaceExW, GetDriveTypeW, GetLogicalDrives, GetVolumeInformationW,
    OPEN_EXISTING,
};
use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
use winapi::um::ioapiset::DeviceIoControl;
use winapi::um::winbase::{DRIVE_FIXED, DRIVE_REMOVABLE};
use winapi::um::winioctl::{
    DEVICE_TRIM_DESCRIPTOR, IOCTL_STORAGE_QUERY_PROPERTY, STORAGE_PROPERTY_QUERY,
};
use winapi::um::winnt::{FILE_SHARE_READ, FILE_SHARE_WRITE, HANDLE, ULARGE_INTEGER};

#[doc = include_str!("../../md_doc/disk.md")]
pub struct Disk {
    type_: DiskType,
    name: OsString,
    file_system: Vec<u8>,
    mount_point: Vec<u16>,
    s_mount_point: String,
    total_space: u64,
    available_space: u64,
    is_removable: bool,
}

impl DiskExt for Disk {
    fn type_(&self) -> DiskType {
        self.type_
    }

    fn name(&self) -> &OsStr {
        &self.name
    }

    fn file_system(&self) -> &[u8] {
        &self.file_system
    }

    fn mount_point(&self) -> &Path {
        Path::new(&self.s_mount_point)
    }

    fn total_space(&self) -> u64 {
        self.total_space
    }

    fn available_space(&self) -> u64 {
        self.available_space
    }

    fn is_removable(&self) -> bool {
        self.is_removable
    }

    fn refresh(&mut self) -> bool {
        if self.total_space != 0 {
            unsafe {
                let mut tmp: ULARGE_INTEGER = std::mem::zeroed();
                if GetDiskFreeSpaceExW(
                    self.mount_point.as_ptr(),
                    std::ptr::null_mut(),
                    std::ptr::null_mut(),
                    &mut tmp,
                ) != 0
                {
                    self.available_space = *tmp.QuadPart();
                    return true;
                }
            }
        }
        false
    }
}

struct HandleWrapper(HANDLE);

impl HandleWrapper {
    unsafe fn new(drive_name: &[u16], open_rights: DWORD) -> Option<Self> {
        let handle = CreateFileW(
            drive_name.as_ptr(),
            open_rights,
            FILE_SHARE_READ | FILE_SHARE_WRITE,
            std::ptr::null_mut(),
            OPEN_EXISTING,
            0,
            std::ptr::null_mut(),
        );
        if handle == INVALID_HANDLE_VALUE {
            CloseHandle(handle);
            None
        } else {
            Some(Self(handle))
        }
    }
}

impl Drop for HandleWrapper {
    fn drop(&mut self) {
        unsafe {
            CloseHandle(self.0);
        }
    }
}

unsafe fn get_drive_size(mount_point: &[u16]) -> Option<(u64, u64)> {
    let mut total_size: ULARGE_INTEGER = std::mem::zeroed();
    let mut available_space: ULARGE_INTEGER = std::mem::zeroed();
    if GetDiskFreeSpaceExW(
        mount_point.as_ptr(),
        std::ptr::null_mut(),
        &mut total_size,
        &mut available_space,
    ) != 0
    {
        Some((
            *total_size.QuadPart() as _,
            *available_space.QuadPart() as _,
        ))
    } else {
        None
    }
}

pub(crate) unsafe fn get_disks() -> Vec<Disk> {
    let drives = GetLogicalDrives();
    if drives == 0 {
        return Vec::new();
    }

    #[cfg(feature = "multithread")]
    use rayon::iter::ParallelIterator;

    crate::utils::into_iter(0..DWORD::BITS)
        .filter_map(|x| {
            if (drives >> x) & 1 == 0 {
                return None;
            }
            let mount_point = [b'A' as u16 + x as u16, b':' as u16, b'\\' as u16, 0];

            let drive_type = GetDriveTypeW(mount_point.as_ptr());

            let is_removable = drive_type == DRIVE_REMOVABLE;

            if drive_type != DRIVE_FIXED && drive_type != DRIVE_REMOVABLE {
                return None;
            }
            let mut name = [0u16; MAX_PATH + 1];
            let mut file_system = [0u16; 32];
            if GetVolumeInformationW(
                mount_point.as_ptr(),
                name.as_mut_ptr(),
                name.len() as DWORD,
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                file_system.as_mut_ptr(),
                file_system.len() as DWORD,
            ) == 0
            {
                return None;
            }
            let mut pos = 0;
            for x in name.iter() {
                if *x == 0 {
                    break;
                }
                pos += 1;
            }
            let name = String::from_utf16_lossy(&name[..pos]);
            let name = OsStr::new(&name);

            pos = 0;
            for x in file_system.iter() {
                if *x == 0 {
                    break;
                }
                pos += 1;
            }
            let file_system: Vec<u8> = file_system[..pos].iter().map(|x| *x as u8).collect();

            let drive_name = [
                b'\\' as u16,
                b'\\' as u16,
                b'.' as u16,
                b'\\' as u16,
                b'A' as u16 + x as u16,
                b':' as u16,
                0,
            ];
            let handle = HandleWrapper::new(&drive_name, 0)?;
            let (total_space, available_space) = get_drive_size(&mount_point)?;
            if total_space == 0 {
                return None;
            }
            /*let mut spq_trim: STORAGE_PROPERTY_QUERY = std::mem::zeroed();
            spq_trim.PropertyId = StorageDeviceTrimProperty;
            spq_trim.QueryType = PropertyStandardQuery;
            let mut dtd: DEVICE_TRIM_DESCRIPTOR = std::mem::zeroed();*/
            let mut spq_trim = STORAGE_PROPERTY_QUERY {
                PropertyId: 8,
                QueryType: 0,
                AdditionalParameters: [0],
            };
            let mut dtd: DEVICE_TRIM_DESCRIPTOR = std::mem::zeroed();

            let mut dw_size = 0;
            let type_ = if DeviceIoControl(
                handle.0,
                IOCTL_STORAGE_QUERY_PROPERTY,
                &mut spq_trim as *mut STORAGE_PROPERTY_QUERY as *mut c_void,
                size_of::<STORAGE_PROPERTY_QUERY>() as DWORD,
                &mut dtd as *mut DEVICE_TRIM_DESCRIPTOR as *mut c_void,
                size_of::<DEVICE_TRIM_DESCRIPTOR>() as DWORD,
                &mut dw_size,
                std::ptr::null_mut(),
            ) == 0
                || dw_size != size_of::<DEVICE_TRIM_DESCRIPTOR>() as DWORD
            {
                DiskType::Unknown(-1)
            } else {
                let is_ssd = dtd.TrimEnabled != 0;
                if is_ssd {
                    DiskType::SSD
                } else {
                    DiskType::HDD
                }
            };
            Some(Disk {
                type_,
                name: name.to_owned(),
                file_system: file_system.to_vec(),
                mount_point: mount_point.to_vec(),
                s_mount_point: String::from_utf16_lossy(&mount_point[..mount_point.len() - 1]),
                total_space,
                available_space,
                is_removable,
            })
        })
        .collect::<Vec<_>>()
}
<commit_msg>Fix Disk kind information retrieval<commit_after>// Take a look at the license at the top of the repository in the LICENSE file.

use crate::{DiskExt, DiskType};

use std::ffi::{OsStr, OsString};
use std::mem::size_of;
use std::path::Path;

use winapi::ctypes::c_void;
use winapi::shared::minwindef::{DWORD, MAX_PATH};
use winapi::um::fileapi::{
    CreateFileW, GetDiskFreeSpaceExW, GetDriveTypeW, GetLogicalDrives, GetVolumeInformationW,
    OPEN_EXISTING,
};
use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
use winapi::um::ioapiset::DeviceIoControl;
use winapi::um::winbase::{DRIVE_FIXED, DRIVE_REMOVABLE};
use winapi::um::winioctl::{
    PropertyStandardQuery, StorageDeviceSeekPenaltyProperty, IOCTL_STORAGE_QUERY_PROPERTY,
    STORAGE_PROPERTY_QUERY,
};
use winapi::um::winnt::{BOOLEAN, FILE_SHARE_READ, FILE_SHARE_WRITE, HANDLE, ULARGE_INTEGER};

#[doc = include_str!("../../md_doc/disk.md")]
pub struct Disk {
    type_: DiskType,
    name: OsString,
    file_system: Vec<u8>,
    mount_point: Vec<u16>,
    s_mount_point: String,
    total_space: u64,
    available_space: u64,
    is_removable: bool,
}

impl DiskExt for Disk {
    fn type_(&self) -> DiskType {
        self.type_
    }

    fn name(&self) -> &OsStr {
        &self.name
    }

    fn file_system(&self) -> &[u8] {
        &self.file_system
    }

    fn mount_point(&self) -> &Path {
        Path::new(&self.s_mount_point)
    }

    fn total_space(&self) -> u64 {
        self.total_space
    }

    fn available_space(&self) -> u64 {
        self.available_space
    }

    fn is_removable(&self) -> bool {
        self.is_removable
    }

    fn refresh(&mut self) -> bool {
        if self.total_space != 0 {
            unsafe {
                let mut tmp: ULARGE_INTEGER = std::mem::zeroed();
                if GetDiskFreeSpaceExW(
                    self.mount_point.as_ptr(),
                    std::ptr::null_mut(),
                    std::ptr::null_mut(),
                    &mut tmp,
                ) != 0
                {
                    self.available_space = *tmp.QuadPart();
                    return true;
                }
            }
        }
        false
    }
}

struct HandleWrapper(HANDLE);

impl HandleWrapper {
    unsafe fn new(drive_name: &[u16], open_rights: DWORD) -> Option<Self> {
        let handle = CreateFileW(
            drive_name.as_ptr(),
            open_rights,
            FILE_SHARE_READ | FILE_SHARE_WRITE,
            std::ptr::null_mut(),
            OPEN_EXISTING,
            0,
            std::ptr::null_mut(),
        );
        if handle == INVALID_HANDLE_VALUE {
            CloseHandle(handle);
            None
        } else {
            Some(Self(handle))
        }
    }
}

impl Drop for HandleWrapper {
    fn drop(&mut self) {
        unsafe {
            CloseHandle(self.0);
        }
    }
}

unsafe fn get_drive_size(mount_point: &[u16]) -> Option<(u64, u64)> {
    let mut total_size: ULARGE_INTEGER = std::mem::zeroed();
    let mut available_space: ULARGE_INTEGER = std::mem::zeroed();
    if GetDiskFreeSpaceExW(
        mount_point.as_ptr(),
        std::ptr::null_mut(),
        &mut total_size,
        &mut available_space,
    ) != 0
    {
        Some((
            *total_size.QuadPart() as _,
            *available_space.QuadPart() as _,
        ))
    } else {
        None
    }
}

// FIXME: To be removed once <https://github.com/retep998/winapi-rs/pull/1028> has been merged.
#[allow(non_snake_case)]
#[repr(C)]
struct DEVICE_SEEK_PENALTY_DESCRIPTOR {
    Version: DWORD,
    Size: DWORD,
    IncursSeekPenalty: BOOLEAN,
}

pub(crate) unsafe fn get_disks() -> Vec<Disk> {
    let drives = GetLogicalDrives();
    if drives == 0 {
        return Vec::new();
    }

    #[cfg(feature = "multithread")]
    use rayon::iter::ParallelIterator;

    crate::utils::into_iter(0..DWORD::BITS)
        .filter_map(|x| {
            if (drives >> x) & 1 == 0 {
                return None;
            }
            let mount_point = [b'A' as u16 + x as u16, b':' as u16, b'\\' as u16, 0];

            let drive_type = GetDriveTypeW(mount_point.as_ptr());

            let is_removable = drive_type == DRIVE_REMOVABLE;

            if drive_type != DRIVE_FIXED && drive_type != DRIVE_REMOVABLE {
                return None;
            }
            let mut name = [0u16; MAX_PATH + 1];
            let mut file_system = [0u16; 32];
            if GetVolumeInformationW(
                mount_point.as_ptr(),
                name.as_mut_ptr(),
                name.len() as DWORD,
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                file_system.as_mut_ptr(),
                file_system.len() as DWORD,
            ) == 0
            {
                return None;
            }
            let mut pos = 0;
            for x in name.iter() {
                if *x == 0 {
                    break;
                }
                pos += 1;
            }
            let name = String::from_utf16_lossy(&name[..pos]);
            let name = OsStr::new(&name);

            pos = 0;
            for x in file_system.iter() {
                if *x == 0 {
                    break;
                }
                pos += 1;
            }
            let file_system: Vec<u8> = file_system[..pos].iter().map(|x| *x as u8).collect();

            let drive_name = [
                b'\\' as u16,
                b'\\' as u16,
                b'.' as u16,
                b'\\' as u16,
                b'A' as u16 + x as u16,
                b':' as u16,
                0,
            ];
            let handle = HandleWrapper::new(&drive_name, 0)?;
            let (total_space, available_space) = get_drive_size(&mount_point)?;
            if total_space == 0 {
                return None;
            }
            let mut spq_trim = STORAGE_PROPERTY_QUERY {
                PropertyId: StorageDeviceSeekPenaltyProperty,
                QueryType: PropertyStandardQuery,
                AdditionalParameters: [0],
            };
            let mut result: DEVICE_SEEK_PENALTY_DESCRIPTOR = std::mem::zeroed();

            let mut dw_size = 0;
            let type_ = if DeviceIoControl(
                handle.0,
                IOCTL_STORAGE_QUERY_PROPERTY,
                &mut spq_trim as *mut STORAGE_PROPERTY_QUERY as *mut c_void,
                size_of::<STORAGE_PROPERTY_QUERY>() as DWORD,
                &mut result as *mut DEVICE_SEEK_PENALTY_DESCRIPTOR as *mut c_void,
                size_of::<DEVICE_SEEK_PENALTY_DESCRIPTOR>() as DWORD,
                &mut dw_size,
                std::ptr::null_mut(),
            ) == 0
                || dw_size != size_of::<DEVICE_SEEK_PENALTY_DESCRIPTOR>() as DWORD
            {
                DiskType::Unknown(-1)
            } else {
                let is_ssd = result.IncursSeekPenalty == 0;
                if is_ssd {
                    DiskType::SSD
                } else {
                    DiskType::HDD
                }
            };
            Some(Disk {
                type_,
                name: name.to_owned(),
                file_system: file_system.to_vec(),
                mount_point: mount_point.to_vec(),
                s_mount_point: String::from_utf16_lossy(&mount_point[..mount_point.len() - 1]),
                total_space,
                available_space,
                is_removable,
            })
        })
        .collect::<Vec<_>>()
}
<|endoftext|>


<commit_before>use std::io::prelude::*;
use std::path::{Path};
use std::fs::{self, File};
use rand::StdRng;
use rand::SeedableRng;
use std::iter::Filter;
use calx::Dijkstra;
use super::entity::Entity;
use super::ecs::EntityIter;
use super::world;
use super::flags;
use super::dir6::Dir6;
use super::area::Area;
use super::location::Location;
use super::ecs::{ComponentAccess};
use super::msg;
use super::{Biome, AreaSpec, Msg};

/// Game update control.
#[derive(Copy, Clone, PartialEq)]
pub enum ControlState {
    AwaitingInput,
    ReadyToUpdate,
}

/// Player input action.
#[derive(Copy, Eq, PartialEq, Clone, Debug, RustcEncodable, RustcDecodable)]
pub enum Input {
    /// Take a step in the given direction.
    Step(Dir6),
    /// Melee attack in the given direction.
    Melee(Dir6),
    /// Shoot in the given direction.
    Shoot(Dir6),
    /// Do nothing for a turn.
    Pass,
}

/// Return the player entity if one exists.
pub fn player() -> Option<Entity> {
    world::with(|w| w.flags.player)
}

/// Return true if the game has ended and the player can make no further
/// actions.
pub fn is_game_over() -> bool {
    player().is_none()
}

/// Find the first entity that has a local (not inherited) Desc component with
/// the given name.
pub fn find_prototype(name: &str) -> Option<Entity> {
    world::with(|w|
        entities().find(|&e| {
            if let Some(d) = w.descs().get_local(e) {
                if d.name == name { return true; }
            }
            false
        })
    )
}

/// Spawn a specific type of entity
pub fn spawn_named(name: &str, loc: Location) {
    find_prototype(name).expect(&format!("Spawn prototype '{}' not found", name)[..])
        .clone_at(loc);
}

// World update state machine //////////////////////////////////////////

/// Get the current control state.
pub fn control_state() -> ControlState {
    if world::with(|w| w.flags.player_acted) { return ControlState::ReadyToUpdate; }
    match player() {
        Some(p) if p.acts_this_frame() => ControlState::AwaitingInput,
        _ => ControlState::ReadyToUpdate,
    }
}

/// Top-level game state update function. Only valid to call if
/// control_state() returned ReadyToUpdate.
pub fn update() {
    assert!(control_state() == ControlState::ReadyToUpdate);

    ai_main();

    world::with_mut(|w| {
        w.flags.tick += 1;
        w.flags.player_acted = false;
    });
}

/// Give player input. Only valid to call if control_state() returned
/// AwaitingInput.
pub fn input(input: Input) {
    assert!(control_state() == ControlState::AwaitingInput);
    let p = player().expect("No player to receive input");
    match input {
        Input::Step(d) => {
            p.step(d);
        }
        Input::Melee(d) => {
            p.melee(d);
        }
        Input::Shoot(d) => {
            p.shoot(d);
        }
        Input::Pass => {
        }
    }
    world::with_mut(|w| w.flags.player_acted = true);

    // Run one world update cycle right away, so that we don't get awkward
    // single frames rendered where the player has acted and the rest of the
    // world hasn't.
    if control_state() == ControlState::ReadyToUpdate {
        update();
    }
}

// Entities ////////////////////////////////////////////////////////////

/// Return an iterator of all the world entities.
pub fn entities() -> EntityIter {
    world::with(|w| w.ecs.iter())
}

/// Return an iterator of all the world mobs.
pub fn _mobs() -> Filter<EntityIter, fn(&Entity) -> bool> {
    fn _is_mob(e: &Entity) -> bool { e.is_mob() }
    entities().filter(_is_mob)
}

/// Run AI for all autonomous mobs.
fn ai_main() {
    for entity in entities() {
        entity.update();
    }
}

// World logic /////////////////////////////////////////////////////////

/// Return the current floor depth. Greater depths mean more powerful monsters
/// and stranger terrain.
pub fn current_depth() -> i32 { world::with(|w| w.area.seed.spec.depth) }

pub fn start_level(depth: i32) {
    let biome = match depth {
        1 => Biome::Overland,
        _ => Biome::Dungeon,
    };

    clear_nonplayers();

    let seed = world::with(|w| w.flags.seed);

    let new_area = Area::new(
        seed,
        AreaSpec::new(biome, depth));
    // XXX: How to move area into the closure without cloning?
    world::with_mut(|w| {
        w.area = new_area.clone();
    });

    let mut rng: StdRng = SeedableRng::from_seed(&[seed as usize + depth as usize][..]);
    for (spawn, loc) in world::with(|w| w.area.get_spawns()).into_iter() {
        spawn.spawn(&mut rng, loc);
    }

    let start_loc = world::with(|w| w.area.player_entrance());
    // Either reuse the existing player or create a new one.
    match player() {
        Some(p) => {
            p.forget_map();
            p.place(start_loc);
        }
        None => {
            let player = find_prototype("player").expect("No Player prototype found!")
            .clone_at(start_loc);
            world::with_mut(|w| w.flags.player = Some(player));
        }
    };
    flags::set_camera(start_loc);
}

fn clear_nonplayers() {
    let po = player();
    for e in entities() {
        // Don't destroy player or player's inventory.
        if let Some(p) = po {
            if e == p || p.contains(e) {
                continue;
            }
        }

        if e.location().is_some() {
            e.delete();
        }
    }
}

/// Move the player to the next level.
pub fn next_level() {
    // This is assuming a really simple, original Rogue style descent-only, no
    // persistent maps style world.
    start_level(current_depth() + 1);
    caption!("Depth {}", current_depth() - 1);
}

// Effects /////////////////////////////////////////////////////////////

/// Create a projectile arc in dir from origin.
pub fn shoot(origin: Location, dir: Dir6, range: u32, power: i32) {
    let mut loc = origin;
    if range == 0 { return; }
    for i in 1..(range + 1) {
        loc = origin + dir.to_v2() * (i as i32);
        if loc.terrain().blocks_shot() {
            msg::push(Msg::Sparks(loc));
            break;
        }
        if let Some(e) = loc.mob_at() {
            e.damage(power);
            break;
        }
    }
    msg::push(Msg::Beam(origin, loc));
}

/// Generate an explosion at location.
pub fn explode(center: Location, power: i32) {
    // Add more complex parametrization if needed.
    msg::push(Msg::Explosion(center));

    // Damage the center.
    if let Some(e) = center.mob_at() {
        e.damage(power);
    }

    for d in Dir6::iter() {
        if let Some(e) = (center + d.to_v2()).mob_at() {
            // Explosions with enough power push back mobs.
            // TODO: Pushback might consider mob size, big guys get pushed
            // around less.
            let push_threshold = 4;
            if power >= push_threshold {
                e.push(*d, ((power - push_threshold) / 2 + 1) as u32);
            }
            if e.exists() {
                e.damage(power);
            }
        }
    }
}

////////////////////////////////////////////////////////////////////////

/// Build a Dijkstra map towards the unexplored corners of the player's FOV.
///
/// Pathing_depth is the depth of the search map. Low pathing depths may not
/// reach distant unexplored cells, but high pathing depths take longer to
/// calculate.
pub fn autoexplore_map(pathing_depth: u32) -> Option<Dijkstra<Location>> {
    let locs = world::with(|w| w.area.terrain.iter()
                           .map(|(&loc, _)| loc)
                           .filter(|loc| loc.fov_status().is_none())
                           .collect::<Vec<Location>>());

    if locs.len() == 0 {
        return None;
    }

    Some(Dijkstra::new(locs, |&loc| !loc.blocks_walk(), pathing_depth))
}

/// Look for targets to shoot in a direction.
pub fn find_target(shooter: Entity, dir: Dir6, range: usize) -> Option<Entity> {
    let origin = shooter.location().unwrap();
    for i in 1..(range + 1) {
        let loc = origin + dir.to_v2() * (i as i32);
        if loc.terrain().blocks_shot() {
            break;
        }
        if let Some(e) = loc.mob_at() {
            if shooter.is_hostile_to(e) { return Some(e); }
        }
    }
    None
}

////////////////////////////////////////////////////////////////////////

pub fn save_game() {
    // Only save if there's still a living player around.
    if is_game_over() {
        return;
    }

    let save_data = world::save();
    File::create("magog_save.json").unwrap()
        .write_all(&save_data.into_bytes()).unwrap();
}

pub fn load_game() {
    let path = Path::new("magog_save.json");
    if !path.exists() { return; }
    let mut save_data = String::new();
    File::open(&path).unwrap().read_to_string(&mut save_data).unwrap();
    // TODO: Handle failed load nicely.
    world::load(&save_data[..]).unwrap();
}

pub fn _delete_save() {
    fs::remove_file("magog_save.json").unwrap();
}

pub fn _save_exists() -> bool  { Path::new("magog_save.json").exists() }
<commit_msg>Don't crash when trying to remove nonexistent save<commit_after>use std::io::prelude::*;
use std::path::{Path};
use std::fs::{self, File};
use rand::StdRng;
use rand::SeedableRng;
use std::iter::Filter;
use calx::Dijkstra;
use super::entity::Entity;
use super::ecs::EntityIter;
use super::world;
use super::flags;
use super::dir6::Dir6;
use super::area::Area;
use super::location::Location;
use super::ecs::{ComponentAccess};
use super::msg;
use super::{Biome, AreaSpec, Msg};

/// Game update control.
#[derive(Copy, Clone, PartialEq)]
pub enum ControlState {
    AwaitingInput,
    ReadyToUpdate,
}

/// Player input action.
#[derive(Copy, Eq, PartialEq, Clone, Debug, RustcEncodable, RustcDecodable)]
pub enum Input {
    /// Take a step in the given direction.
    Step(Dir6),
    /// Melee attack in the given direction.
    Melee(Dir6),
    /// Shoot in the given direction.
    Shoot(Dir6),
    /// Do nothing for a turn.
    Pass,
}

/// Return the player entity if one exists.
pub fn player() -> Option<Entity> {
    world::with(|w| w.flags.player)
}

/// Return true if the game has ended and the player can make no further
/// actions.
pub fn is_game_over() -> bool {
    player().is_none()
}

/// Find the first entity that has a local (not inherited) Desc component with
/// the given name.
pub fn find_prototype(name: &str) -> Option<Entity> {
    world::with(|w|
        entities().find(|&e| {
            if let Some(d) = w.descs().get_local(e) {
                if d.name == name { return true; }
            }
            false
        })
    )
}

/// Spawn a specific type of entity
pub fn spawn_named(name: &str, loc: Location) {
    find_prototype(name).expect(&format!("Spawn prototype '{}' not found", name)[..])
        .clone_at(loc);
}

// World update state machine //////////////////////////////////////////

/// Get the current control state.
pub fn control_state() -> ControlState {
    if world::with(|w| w.flags.player_acted) { return ControlState::ReadyToUpdate; }
    match player() {
        Some(p) if p.acts_this_frame() => ControlState::AwaitingInput,
        _ => ControlState::ReadyToUpdate,
    }
}

/// Top-level game state update function. Only valid to call if
/// control_state() returned ReadyToUpdate.
pub fn update() {
    assert!(control_state() == ControlState::ReadyToUpdate);

    ai_main();

    world::with_mut(|w| {
        w.flags.tick += 1;
        w.flags.player_acted = false;
    });
}

/// Give player input. Only valid to call if control_state() returned
/// AwaitingInput.
pub fn input(input: Input) {
    assert!(control_state() == ControlState::AwaitingInput);
    let p = player().expect("No player to receive input");
    match input {
        Input::Step(d) => {
            p.step(d);
        }
        Input::Melee(d) => {
            p.melee(d);
        }
        Input::Shoot(d) => {
            p.shoot(d);
        }
        Input::Pass => {
        }
    }
    world::with_mut(|w| w.flags.player_acted = true);

    // Run one world update cycle right away, so that we don't get awkward
    // single frames rendered where the player has acted and the rest of the
    // world hasn't.
    if control_state() == ControlState::ReadyToUpdate {
        update();
    }
}

// Entities ////////////////////////////////////////////////////////////

/// Return an iterator of all the world entities.
pub fn entities() -> EntityIter {
    world::with(|w| w.ecs.iter())
}

/// Return an iterator of all the world mobs.
pub fn _mobs() -> Filter<EntityIter, fn(&Entity) -> bool> {
    fn _is_mob(e: &Entity) -> bool { e.is_mob() }
    entities().filter(_is_mob)
}

/// Run AI for all autonomous mobs.
fn ai_main() {
    for entity in entities() {
        entity.update();
    }
}

// World logic /////////////////////////////////////////////////////////

/// Return the current floor depth. Greater depths mean more powerful monsters
/// and stranger terrain.
pub fn current_depth() -> i32 { world::with(|w| w.area.seed.spec.depth) }

pub fn start_level(depth: i32) {
    let biome = match depth {
        1 => Biome::Overland,
        _ => Biome::Dungeon,
    };

    clear_nonplayers();

    let seed = world::with(|w| w.flags.seed);

    let new_area = Area::new(
        seed,
        AreaSpec::new(biome, depth));
    // XXX: How to move area into the closure without cloning?
    world::with_mut(|w| {
        w.area = new_area.clone();
    });

    let mut rng: StdRng = SeedableRng::from_seed(&[seed as usize + depth as usize][..]);
    for (spawn, loc) in world::with(|w| w.area.get_spawns()).into_iter() {
        spawn.spawn(&mut rng, loc);
    }

    let start_loc = world::with(|w| w.area.player_entrance());
    // Either reuse the existing player or create a new one.
    match player() {
        Some(p) => {
            p.forget_map();
            p.place(start_loc);
        }
        None => {
            let player = find_prototype("player").expect("No Player prototype found!")
            .clone_at(start_loc);
            world::with_mut(|w| w.flags.player = Some(player));
        }
    };
    flags::set_camera(start_loc);
}

fn clear_nonplayers() {
    let po = player();
    for e in entities() {
        // Don't destroy player or player's inventory.
        if let Some(p) = po {
            if e == p || p.contains(e) {
                continue;
            }
        }

        if e.location().is_some() {
            e.delete();
        }
    }
}

/// Move the player to the next level.
pub fn next_level() {
    // This is assuming a really simple, original Rogue style descent-only, no
    // persistent maps style world.
    start_level(current_depth() + 1);
    caption!("Depth {}", current_depth() - 1);
}

// Effects /////////////////////////////////////////////////////////////

/// Create a projectile arc in dir from origin.
pub fn shoot(origin: Location, dir: Dir6, range: u32, power: i32) {
    let mut loc = origin;
    if range == 0 { return; }
    for i in 1..(range + 1) {
        loc = origin + dir.to_v2() * (i as i32);
        if loc.terrain().blocks_shot() {
            msg::push(Msg::Sparks(loc));
            break;
        }
        if let Some(e) = loc.mob_at() {
            e.damage(power);
            break;
        }
    }
    msg::push(Msg::Beam(origin, loc));
}

/// Generate an explosion at location.
pub fn explode(center: Location, power: i32) {
    // Add more complex parametrization if needed.
    msg::push(Msg::Explosion(center));

    // Damage the center.
    if let Some(e) = center.mob_at() {
        e.damage(power);
    }

    for d in Dir6::iter() {
        if let Some(e) = (center + d.to_v2()).mob_at() {
            // Explosions with enough power push back mobs.
            // TODO: Pushback might consider mob size, big guys get pushed
            // around less.
            let push_threshold = 4;
            if power >= push_threshold {
                e.push(*d, ((power - push_threshold) / 2 + 1) as u32);
            }
            if e.exists() {
                e.damage(power);
            }
        }
    }
}

////////////////////////////////////////////////////////////////////////

/// Build a Dijkstra map towards the unexplored corners of the player's FOV.
///
/// Pathing_depth is the depth of the search map. Low pathing depths may not
/// reach distant unexplored cells, but high pathing depths take longer to
/// calculate.
pub fn autoexplore_map(pathing_depth: u32) -> Option<Dijkstra<Location>> {
    let locs = world::with(|w| w.area.terrain.iter()
                           .map(|(&loc, _)| loc)
                           .filter(|loc| loc.fov_status().is_none())
                           .collect::<Vec<Location>>());

    if locs.len() == 0 {
        return None;
    }

    Some(Dijkstra::new(locs, |&loc| !loc.blocks_walk(), pathing_depth))
}

/// Look for targets to shoot in a direction.
pub fn find_target(shooter: Entity, dir: Dir6, range: usize) -> Option<Entity> {
    let origin = shooter.location().unwrap();
    for i in 1..(range + 1) {
        let loc = origin + dir.to_v2() * (i as i32);
        if loc.terrain().blocks_shot() {
            break;
        }
        if let Some(e) = loc.mob_at() {
            if shooter.is_hostile_to(e) { return Some(e); }
        }
    }
    None
}

////////////////////////////////////////////////////////////////////////

pub fn save_game() {
    // Only save if there's still a living player around.
    if is_game_over() {
        return;
    }

    let save_data = world::save();
    File::create("magog_save.json").unwrap()
        .write_all(&save_data.into_bytes()).unwrap();
}

pub fn load_game() {
    let path = Path::new("magog_save.json");
    if !path.exists() { return; }
    let mut save_data = String::new();
    File::open(&path).unwrap().read_to_string(&mut save_data).unwrap();
    // TODO: Handle failed load nicely.
    world::load(&save_data[..]).unwrap();
}

pub fn _delete_save() {
    let _ = fs::remove_file("magog_save.json");
}

pub fn _save_exists() -> bool  { Path::new("magog_save.json").exists() }
<|endoftext|>