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

//! A work queue for scheduling units of work across threads in a fork-join fashion.
//!
//! Data associated with queues is simply a pair of unsigned integers. It is expected that a
//! higher-level API on top of this could allow safe fork-join parallelism.

#[cfg(windows)]
extern crate kernel32;

use deque::{self, Abort, Data, Empty, Stealer, Worker};
#[cfg(not(windows))]
use libc::usleep;
use rand::{Rng, XorShiftRng, weak_rng};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{Receiver, Sender, channel};
use thread::spawn_named;
use thread_state;

/// A unit of work.
///
/// # Type parameters
///
/// - `QueueData`: global custom data for the entire work queue.
/// - `WorkData`: custom data specific to each unit of work.
pub struct WorkUnit<QueueData, WorkData: Send> {
    /// The function to execute.
    pub fun: extern "Rust" fn(WorkData, &mut WorkerProxy<QueueData, WorkData>),
    /// Arbitrary data.
    pub data: WorkData,
}

/// Messages from the supervisor to the worker.
enum WorkerMsg<QueueData: 'static, WorkData: 'static + Send> {
    /// Tells the worker to start work.
    Start(Worker<WorkUnit<QueueData, WorkData>>, *mut AtomicUsize, *const QueueData),
    /// Tells the worker to stop. It can be restarted again with a `WorkerMsg::Start`.
    Stop,
    /// Tells the worker to measure the heap size of its TLS using the supplied function.
    HeapSizeOfTLS(fn() -> usize),
    /// Tells the worker thread to terminate.
    Exit,
}

unsafe impl<QueueData: 'static, WorkData: 'static + Send> Send for WorkerMsg<QueueData, WorkData> {}

/// Messages to the supervisor.
enum SupervisorMsg<QueueData: 'static, WorkData: 'static + Send> {
    Finished,
    HeapSizeOfTLS(usize),
    ReturnDeque(usize, Worker<WorkUnit<QueueData, WorkData>>),
}

unsafe impl<QueueData: 'static, WorkData: 'static + Send> Send for SupervisorMsg<QueueData, WorkData> {}

/// Information that the supervisor thread keeps about the worker threads.
struct WorkerInfo<QueueData: 'static, WorkData: 'static + Send> {
    /// The communication channel to the workers.
    chan: Sender<WorkerMsg<QueueData, WorkData>>,
    /// The worker end of the deque, if we have it.
    deque: Option<Worker<WorkUnit<QueueData, WorkData>>>,
    /// The thief end of the work-stealing deque.
    thief: Stealer<WorkUnit<QueueData, WorkData>>,
}

/// Information specific to each worker thread that the thread keeps.
struct WorkerThread<QueueData: 'static, WorkData: 'static + Send> {
    /// The index of this worker.
    index: usize,
    /// The communication port from the supervisor.
    port: Receiver<WorkerMsg<QueueData, WorkData>>,
    /// The communication channel on which messages are sent to the supervisor.
    chan: Sender<SupervisorMsg<QueueData, WorkData>>,
    /// The thief end of the work-stealing deque for all other workers.
    other_deques: Vec<Stealer<WorkUnit<QueueData, WorkData>>>,
    /// The random number generator for this worker.
    rng: XorShiftRng,
}

unsafe impl<QueueData: 'static, WorkData: 'static + Send> Send for WorkerThread<QueueData, WorkData> {}

const SPINS_UNTIL_BACKOFF: u32 = 128;
const BACKOFF_INCREMENT_IN_US: u32 = 5;
const BACKOFFS_UNTIL_CONTROL_CHECK: u32 = 6;

#[cfg(not(windows))]
fn sleep_microseconds(usec: u32) {
    unsafe {
        usleep(usec);
    }
}

#[cfg(windows)]
fn sleep_microseconds(_: u32) {
    unsafe {
        kernel32::Sleep(0);
    }
}

impl<QueueData: Sync, WorkData: Send> WorkerThread<QueueData, WorkData> {
    /// The main logic. This function starts up the worker and listens for
    /// messages.
    fn start(&mut self) {
        let deque_index_mask = (self.other_deques.len() as u32).next_power_of_two() - 1;
        loop {
            // Wait for a start message.
            let (mut deque, ref_count, queue_data) = match self.port.recv().unwrap() {
                WorkerMsg::Start(deque, ref_count, queue_data) => (deque, ref_count, queue_data),
                WorkerMsg::Stop => panic!("unexpected stop message"),
                WorkerMsg::Exit => return,
                WorkerMsg::HeapSizeOfTLS(f) => {
                    self.chan.send(SupervisorMsg::HeapSizeOfTLS(f())).unwrap();
                    continue;
                }
            };

            let mut back_off_sleep = 0 as u32;

            // We're off!
            'outer: loop {
                let work_unit;
                match deque.pop() {
                    Some(work) => work_unit = work,
                    None => {
                        // Become a thief.
                        let mut i = 0;
                        loop {
                            // Don't just use `rand % len` because that's slow on ARM.
                            let mut victim;
                            loop {
                                victim = self.rng.next_u32() & deque_index_mask;
                                if (victim as usize) < self.other_deques.len() {
                                    break
                                }
                            }

                            match self.other_deques[victim as usize].steal() {
                                Empty | Abort => {
                                    // Continue.
                                }
                                Data(work) => {
                                    work_unit = work;
                                    back_off_sleep = 0 as u32;
                                    break
                                }
                            }

                            if i > SPINS_UNTIL_BACKOFF {
                                if back_off_sleep >= BACKOFF_INCREMENT_IN_US *
                                        BACKOFFS_UNTIL_CONTROL_CHECK {
                                    match self.port.try_recv() {
                                        Ok(WorkerMsg::Stop) => break 'outer,
                                        Ok(WorkerMsg::Exit) => return,
                                        Ok(_) => panic!("unexpected message"),
                                        _ => {}
                                    }
                                }

                                sleep_microseconds(back_off_sleep);

                                back_off_sleep += BACKOFF_INCREMENT_IN_US;
                                i = 0
                            } else {
                                i += 1
                            }
                        }
                    }
                }

                // At this point, we have some work. Perform it.
                let mut proxy = WorkerProxy {
                    worker: &mut deque,
                    ref_count: ref_count,
                    // queue_data is kept alive in the stack frame of
                    // WorkQueue::run until we send the
                    // SupervisorMsg::ReturnDeque message below.
                    queue_data: unsafe { &*queue_data },
                    worker_index: self.index as u8,
                };
                (work_unit.fun)(work_unit.data, &mut proxy);

                // The work is done. Now decrement the count of outstanding work items. If this was
                // the last work unit in the queue, then send a message on the channel.
                unsafe {
                    if (*ref_count).fetch_sub(1, Ordering::Release) == 1 {
                        self.chan.send(SupervisorMsg::Finished).unwrap()
                    }
                }
            }

            // Give the deque back to the supervisor.
            self.chan.send(SupervisorMsg::ReturnDeque(self.index, deque)).unwrap()
        }
    }
}

/// A handle to the work queue that individual work units have.
pub struct WorkerProxy<'a, QueueData: 'a, WorkData: 'a + Send> {
    worker: &'a mut Worker<WorkUnit<QueueData, WorkData>>,
    ref_count: *mut AtomicUsize,
    queue_data: &'a QueueData,
    worker_index: u8,
}

impl<'a, QueueData: 'static, WorkData: Send + 'static> WorkerProxy<'a, QueueData, WorkData> {
    /// Enqueues a block into the work queue.
    #[inline]
    pub fn push(&mut self, work_unit: WorkUnit<QueueData, WorkData>) {
        unsafe {
            drop((*self.ref_count).fetch_add(1, Ordering::Relaxed));
        }
        self.worker.push(work_unit);
    }

    /// Retrieves the queue user data.
    #[inline]
    pub fn user_data(&self) -> &'a QueueData {
        self.queue_data
    }

    /// Retrieves the index of the worker.
    #[inline]
    pub fn worker_index(&self) -> u8 {
        self.worker_index
    }
}

/// A work queue on which units of work can be submitted.
pub struct WorkQueue<QueueData: 'static, WorkData: 'static + Send> {
    /// Information about each of the workers.
    workers: Vec<WorkerInfo<QueueData, WorkData>>,
    /// A port on which deques can be received from the workers.
    port: Receiver<SupervisorMsg<QueueData, WorkData>>,
    /// The amount of work that has been enqueued.
    work_count: usize,
}

impl<QueueData: Sync, WorkData: Send> WorkQueue<QueueData, WorkData> {
    /// Creates a new work queue and spawns all the threads associated with
    /// it.
    pub fn new(thread_name: &'static str,
               state: thread_state::ThreadState,
               thread_count: usize) -> WorkQueue<QueueData, WorkData> {
        // Set up data structures.
        let (supervisor_chan, supervisor_port) = channel();
        let (mut infos, mut threads) = (vec!(), vec!());
        for i in 0..thread_count {
            let (worker_chan, worker_port) = channel();
            let (worker, thief) = deque::new();
            infos.push(WorkerInfo {
                chan: worker_chan,
                deque: Some(worker),
                thief: thief,
            });
            threads.push(WorkerThread {
                index: i,
                port: worker_port,
                chan: supervisor_chan.clone(),
                other_deques: vec!(),
                rng: weak_rng(),
            });
        }

        // Connect workers to one another.
        for (i, mut thread) in threads.iter_mut().enumerate() {
            for (j, info) in infos.iter().enumerate() {
                if i != j {
                    thread.other_deques.push(info.thief.clone())
                }
            }
            assert!(thread.other_deques.len() == thread_count - 1)
        }

        // Spawn threads.
        for (i, thread) in threads.into_iter().enumerate() {
            spawn_named(
                format!("{} worker {}/{}", thread_name, i + 1, thread_count),
                move || {
                    thread_state::initialize(state | thread_state::IN_WORKER);
                    let mut thread = thread;
                    thread.start()
                })
        }

        WorkQueue {
            workers: infos,
            port: supervisor_port,
            work_count: 0,
        }
    }

    /// Enqueues a block into the work queue.
    #[inline]
    pub fn push(&mut self, work_unit: WorkUnit<QueueData, WorkData>) {
        let deque = &mut self.workers[0].deque;
        match *deque {
            None => {
                panic!("tried to push a block but we don't have the deque?!")
            }
            Some(ref mut deque) => deque.push(work_unit),
        }
        self.work_count += 1
    }

    /// Synchronously runs all the enqueued tasks and waits for them to complete.
    pub fn run(&mut self, data: &QueueData) {
        // Tell the workers to start.
        let mut work_count = AtomicUsize::new(self.work_count);
        for worker in &mut self.workers {
            worker.chan.send(WorkerMsg::Start(worker.deque.take().unwrap(),
                                              &mut work_count,
                                              data)).unwrap()
        }

        // Wait for the work to finish.
        drop(self.port.recv());
        self.work_count = 0;

        // Tell everyone to stop.
        for worker in &self.workers {
            worker.chan.send(WorkerMsg::Stop).unwrap()
        }

        // Get our deques back.
        for _ in 0..self.workers.len() {
            match self.port.recv().unwrap() {
                SupervisorMsg::ReturnDeque(index, deque) => self.workers[index].deque = Some(deque),
                SupervisorMsg::HeapSizeOfTLS(_) => panic!("unexpected HeapSizeOfTLS message"),
                SupervisorMsg::Finished => panic!("unexpected finished message!"),
            }
        }
    }

    /// Synchronously measure memory usage of any thread-local storage.
    pub fn heap_size_of_tls(&self, f: fn() -> usize) -> Vec<usize> {
        // Tell the workers to measure themselves.
        for worker in &self.workers {
            worker.chan.send(WorkerMsg::HeapSizeOfTLS(f)).unwrap()
        }

        // Wait for the workers to finish measuring themselves.
        let mut sizes = vec![];
        for _ in 0..self.workers.len() {
            match self.port.recv().unwrap() {
                SupervisorMsg::HeapSizeOfTLS(size) => {
                    sizes.push(size);
                }
                _ => panic!("unexpected message!"),
            }
        }
        sizes
    }

    pub fn shutdown(&mut self) {
        for worker in &self.workers {
            worker.chan.send(WorkerMsg::Exit).unwrap()
        }
    }
}
<commit_msg>Auto merge of #11797 - Ms2ger:workqueue-mut-atomic, r=nox<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/. */

//! A work queue for scheduling units of work across threads in a fork-join fashion.
//!
//! Data associated with queues is simply a pair of unsigned integers. It is expected that a
//! higher-level API on top of this could allow safe fork-join parallelism.

#[cfg(windows)]
extern crate kernel32;

use deque::{self, Abort, Data, Empty, Stealer, Worker};
#[cfg(not(windows))]
use libc::usleep;
use rand::{Rng, XorShiftRng, weak_rng};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{Receiver, Sender, channel};
use thread::spawn_named;
use thread_state;

/// A unit of work.
///
/// # Type parameters
///
/// - `QueueData`: global custom data for the entire work queue.
/// - `WorkData`: custom data specific to each unit of work.
pub struct WorkUnit<QueueData, WorkData: Send> {
    /// The function to execute.
    pub fun: extern "Rust" fn(WorkData, &mut WorkerProxy<QueueData, WorkData>),
    /// Arbitrary data.
    pub data: WorkData,
}

/// Messages from the supervisor to the worker.
enum WorkerMsg<QueueData: 'static, WorkData: 'static + Send> {
    /// Tells the worker to start work.
    Start(Worker<WorkUnit<QueueData, WorkData>>, *const AtomicUsize, *const QueueData),
    /// Tells the worker to stop. It can be restarted again with a `WorkerMsg::Start`.
    Stop,
    /// Tells the worker to measure the heap size of its TLS using the supplied function.
    HeapSizeOfTLS(fn() -> usize),
    /// Tells the worker thread to terminate.
    Exit,
}

unsafe impl<QueueData: 'static, WorkData: 'static + Send> Send for WorkerMsg<QueueData, WorkData> {}

/// Messages to the supervisor.
enum SupervisorMsg<QueueData: 'static, WorkData: 'static + Send> {
    Finished,
    HeapSizeOfTLS(usize),
    ReturnDeque(usize, Worker<WorkUnit<QueueData, WorkData>>),
}

unsafe impl<QueueData: 'static, WorkData: 'static + Send> Send for SupervisorMsg<QueueData, WorkData> {}

/// Information that the supervisor thread keeps about the worker threads.
struct WorkerInfo<QueueData: 'static, WorkData: 'static + Send> {
    /// The communication channel to the workers.
    chan: Sender<WorkerMsg<QueueData, WorkData>>,
    /// The worker end of the deque, if we have it.
    deque: Option<Worker<WorkUnit<QueueData, WorkData>>>,
    /// The thief end of the work-stealing deque.
    thief: Stealer<WorkUnit<QueueData, WorkData>>,
}

/// Information specific to each worker thread that the thread keeps.
struct WorkerThread<QueueData: 'static, WorkData: 'static + Send> {
    /// The index of this worker.
    index: usize,
    /// The communication port from the supervisor.
    port: Receiver<WorkerMsg<QueueData, WorkData>>,
    /// The communication channel on which messages are sent to the supervisor.
    chan: Sender<SupervisorMsg<QueueData, WorkData>>,
    /// The thief end of the work-stealing deque for all other workers.
    other_deques: Vec<Stealer<WorkUnit<QueueData, WorkData>>>,
    /// The random number generator for this worker.
    rng: XorShiftRng,
}

unsafe impl<QueueData: 'static, WorkData: 'static + Send> Send for WorkerThread<QueueData, WorkData> {}

const SPINS_UNTIL_BACKOFF: u32 = 128;
const BACKOFF_INCREMENT_IN_US: u32 = 5;
const BACKOFFS_UNTIL_CONTROL_CHECK: u32 = 6;

#[cfg(not(windows))]
fn sleep_microseconds(usec: u32) {
    unsafe {
        usleep(usec);
    }
}

#[cfg(windows)]
fn sleep_microseconds(_: u32) {
    unsafe {
        kernel32::Sleep(0);
    }
}

impl<QueueData: Sync, WorkData: Send> WorkerThread<QueueData, WorkData> {
    /// The main logic. This function starts up the worker and listens for
    /// messages.
    fn start(&mut self) {
        let deque_index_mask = (self.other_deques.len() as u32).next_power_of_two() - 1;
        loop {
            // Wait for a start message.
            let (mut deque, ref_count, queue_data) = match self.port.recv().unwrap() {
                WorkerMsg::Start(deque, ref_count, queue_data) => (deque, ref_count, queue_data),
                WorkerMsg::Stop => panic!("unexpected stop message"),
                WorkerMsg::Exit => return,
                WorkerMsg::HeapSizeOfTLS(f) => {
                    self.chan.send(SupervisorMsg::HeapSizeOfTLS(f())).unwrap();
                    continue;
                }
            };

            let mut back_off_sleep = 0 as u32;

            // We're off!
            'outer: loop {
                let work_unit;
                match deque.pop() {
                    Some(work) => work_unit = work,
                    None => {
                        // Become a thief.
                        let mut i = 0;
                        loop {
                            // Don't just use `rand % len` because that's slow on ARM.
                            let mut victim;
                            loop {
                                victim = self.rng.next_u32() & deque_index_mask;
                                if (victim as usize) < self.other_deques.len() {
                                    break
                                }
                            }

                            match self.other_deques[victim as usize].steal() {
                                Empty | Abort => {
                                    // Continue.
                                }
                                Data(work) => {
                                    work_unit = work;
                                    back_off_sleep = 0 as u32;
                                    break
                                }
                            }

                            if i > SPINS_UNTIL_BACKOFF {
                                if back_off_sleep >= BACKOFF_INCREMENT_IN_US *
                                        BACKOFFS_UNTIL_CONTROL_CHECK {
                                    match self.port.try_recv() {
                                        Ok(WorkerMsg::Stop) => break 'outer,
                                        Ok(WorkerMsg::Exit) => return,
                                        Ok(_) => panic!("unexpected message"),
                                        _ => {}
                                    }
                                }

                                sleep_microseconds(back_off_sleep);

                                back_off_sleep += BACKOFF_INCREMENT_IN_US;
                                i = 0
                            } else {
                                i += 1
                            }
                        }
                    }
                }

                // At this point, we have some work. Perform it.
                let mut proxy = WorkerProxy {
                    worker: &mut deque,
                    ref_count: ref_count,
                    // queue_data is kept alive in the stack frame of
                    // WorkQueue::run until we send the
                    // SupervisorMsg::ReturnDeque message below.
                    queue_data: unsafe { &*queue_data },
                    worker_index: self.index as u8,
                };
                (work_unit.fun)(work_unit.data, &mut proxy);

                // The work is done. Now decrement the count of outstanding work items. If this was
                // the last work unit in the queue, then send a message on the channel.
                unsafe {
                    if (*ref_count).fetch_sub(1, Ordering::Release) == 1 {
                        self.chan.send(SupervisorMsg::Finished).unwrap()
                    }
                }
            }

            // Give the deque back to the supervisor.
            self.chan.send(SupervisorMsg::ReturnDeque(self.index, deque)).unwrap()
        }
    }
}

/// A handle to the work queue that individual work units have.
pub struct WorkerProxy<'a, QueueData: 'a, WorkData: 'a + Send> {
    worker: &'a mut Worker<WorkUnit<QueueData, WorkData>>,
    ref_count: *const AtomicUsize,
    queue_data: &'a QueueData,
    worker_index: u8,
}

impl<'a, QueueData: 'static, WorkData: Send + 'static> WorkerProxy<'a, QueueData, WorkData> {
    /// Enqueues a block into the work queue.
    #[inline]
    pub fn push(&mut self, work_unit: WorkUnit<QueueData, WorkData>) {
        unsafe {
            drop((*self.ref_count).fetch_add(1, Ordering::Relaxed));
        }
        self.worker.push(work_unit);
    }

    /// Retrieves the queue user data.
    #[inline]
    pub fn user_data(&self) -> &'a QueueData {
        self.queue_data
    }

    /// Retrieves the index of the worker.
    #[inline]
    pub fn worker_index(&self) -> u8 {
        self.worker_index
    }
}

/// A work queue on which units of work can be submitted.
pub struct WorkQueue<QueueData: 'static, WorkData: 'static + Send> {
    /// Information about each of the workers.
    workers: Vec<WorkerInfo<QueueData, WorkData>>,
    /// A port on which deques can be received from the workers.
    port: Receiver<SupervisorMsg<QueueData, WorkData>>,
    /// The amount of work that has been enqueued.
    work_count: usize,
}

impl<QueueData: Sync, WorkData: Send> WorkQueue<QueueData, WorkData> {
    /// Creates a new work queue and spawns all the threads associated with
    /// it.
    pub fn new(thread_name: &'static str,
               state: thread_state::ThreadState,
               thread_count: usize) -> WorkQueue<QueueData, WorkData> {
        // Set up data structures.
        let (supervisor_chan, supervisor_port) = channel();
        let (mut infos, mut threads) = (vec!(), vec!());
        for i in 0..thread_count {
            let (worker_chan, worker_port) = channel();
            let (worker, thief) = deque::new();
            infos.push(WorkerInfo {
                chan: worker_chan,
                deque: Some(worker),
                thief: thief,
            });
            threads.push(WorkerThread {
                index: i,
                port: worker_port,
                chan: supervisor_chan.clone(),
                other_deques: vec!(),
                rng: weak_rng(),
            });
        }

        // Connect workers to one another.
        for (i, mut thread) in threads.iter_mut().enumerate() {
            for (j, info) in infos.iter().enumerate() {
                if i != j {
                    thread.other_deques.push(info.thief.clone())
                }
            }
            assert!(thread.other_deques.len() == thread_count - 1)
        }

        // Spawn threads.
        for (i, thread) in threads.into_iter().enumerate() {
            spawn_named(
                format!("{} worker {}/{}", thread_name, i + 1, thread_count),
                move || {
                    thread_state::initialize(state | thread_state::IN_WORKER);
                    let mut thread = thread;
                    thread.start()
                })
        }

        WorkQueue {
            workers: infos,
            port: supervisor_port,
            work_count: 0,
        }
    }

    /// Enqueues a block into the work queue.
    #[inline]
    pub fn push(&mut self, work_unit: WorkUnit<QueueData, WorkData>) {
        let deque = &mut self.workers[0].deque;
        match *deque {
            None => {
                panic!("tried to push a block but we don't have the deque?!")
            }
            Some(ref mut deque) => deque.push(work_unit),
        }
        self.work_count += 1
    }

    /// Synchronously runs all the enqueued tasks and waits for them to complete.
    pub fn run(&mut self, data: &QueueData) {
        // Tell the workers to start.
        let work_count = AtomicUsize::new(self.work_count);
        for worker in &mut self.workers {
            worker.chan.send(WorkerMsg::Start(worker.deque.take().unwrap(),
                                              &work_count,
                                              data)).unwrap()
        }

        // Wait for the work to finish.
        drop(self.port.recv());
        self.work_count = 0;

        // Tell everyone to stop.
        for worker in &self.workers {
            worker.chan.send(WorkerMsg::Stop).unwrap()
        }

        // Get our deques back.
        for _ in 0..self.workers.len() {
            match self.port.recv().unwrap() {
                SupervisorMsg::ReturnDeque(index, deque) => self.workers[index].deque = Some(deque),
                SupervisorMsg::HeapSizeOfTLS(_) => panic!("unexpected HeapSizeOfTLS message"),
                SupervisorMsg::Finished => panic!("unexpected finished message!"),
            }
        }
    }

    /// Synchronously measure memory usage of any thread-local storage.
    pub fn heap_size_of_tls(&self, f: fn() -> usize) -> Vec<usize> {
        // Tell the workers to measure themselves.
        for worker in &self.workers {
            worker.chan.send(WorkerMsg::HeapSizeOfTLS(f)).unwrap()
        }

        // Wait for the workers to finish measuring themselves.
        let mut sizes = vec![];
        for _ in 0..self.workers.len() {
            match self.port.recv().unwrap() {
                SupervisorMsg::HeapSizeOfTLS(size) => {
                    sizes.push(size);
                }
                _ => panic!("unexpected message!"),
            }
        }
        sizes
    }

    pub fn shutdown(&mut self) {
        for worker in &self.workers {
            worker.chan.send(WorkerMsg::Exit).unwrap()
        }
    }
}
<|endoftext|>


<commit_before>//! Handle `cargo add` arguments

#![allow(clippy::bool_assert_comparison)]

use cargo_edit::{
    find, get_features_from_registry, get_manifest_from_url, registry_url, workspace_members,
    Dependency,
};
use cargo_edit::{get_latest_dependency, CrateSpec};
use cargo_metadata::Package;
use clap::Parser;
use std::path::PathBuf;

use crate::errors::*;

#[derive(Debug, Parser)]
#[clap(bin_name = "cargo")]
pub enum Command {
    /// Add dependencies to a Cargo.toml manifest file.
    #[clap(name = "add")]
    #[clap(after_help = "\
Examples:
  $ cargo add regex
  $ cargo add regex@0.1.41 --build
  $ cargo add trycmd --dev
  $ cargo add ./crate/parser/
")]
    Add(Args),
}

#[derive(Debug, Parser)]
#[clap(about, version)]
#[clap(setting = clap::AppSettings::DeriveDisplayOrder)]
pub struct Args {
    /// Reference to a package to add as a dependency
    ///
    /// You can reference a packages by:{n}
    /// - `<name>`, like `cargo add serde` (latest version will be used){n}
    /// - `<name>@<version-req>`, like `cargo add serde@1` or `cargo add serde@=1.0.38`{n}
    /// - `<path>`, like `cargo add ./crates/parser/`
    #[clap(value_name = "DEP_ID", required = true)]
    pub crates: Vec<String>,

    /// Disable the default features
    #[clap(long)]
    no_default_features: bool,
    /// Re-enable the default features
    #[clap(long, overrides_with = "no-default-features")]
    default_features: bool,

    /// Space-separated list of features to add
    #[clap(long)]
    pub features: Option<Vec<String>>,

    /// Mark the dependency as optional
    ///
    /// The package name will be exposed as feature of your crate.
    #[clap(long, conflicts_with = "dev")]
    pub optional: bool,

    /// Mark the dependency as required
    ///
    /// The package will be removed from your features.
    #[clap(long, conflicts_with = "dev", overrides_with = "optional")]
    pub no_optional: bool,

    /// Rename the dependency
    ///
    /// Example uses:{n}
    /// - Depending on multiple versions of a crate{n}
    /// - Depend on crates with the same name from different registries
    #[clap(long, short)]
    pub rename: Option<String>,

    /// Package registry for this dependency
    #[clap(long, conflicts_with = "git")]
    pub registry: Option<String>,

    /// Add as development dependency
    ///
    /// Dev-dependencies are not used when compiling a package for building, but are used for compiling tests, examples, and benchmarks.
    ///
    /// These dependencies are not propagated to other packages which depend on this package.
    #[clap(short = 'D', long, help_heading = "SECTION", group = "section")]
    pub dev: bool,

    /// Add as build dependency
    ///
    /// Build-dependencies are the only dependencies available for use by build scripts (`build.rs`
    /// files).
    #[clap(short = 'B', long, help_heading = "SECTION", group = "section")]
    pub build: bool,

    /// Add as dependency to the given target platform.
    #[clap(
        long,
        forbid_empty_values = true,
        help_heading = "SECTION",
        group = "section"
    )]
    pub target: Option<String>,

    /// Path to `Cargo.toml`
    #[clap(long, value_name = "PATH", parse(from_os_str))]
    pub manifest_path: Option<PathBuf>,

    /// Package to modify
    #[clap(short = 'p', long = "package", value_name = "PKGID")]
    pub pkgid: Option<String>,

    /// Run without accessing the network
    #[clap(long)]
    pub offline: bool,

    /// Do not print any output in case of success.
    #[clap(long)]
    pub quiet: bool,

    /// Unstable (nightly-only) flags
    #[clap(
        short = 'Z',
        value_name = "FLAG",
        help_heading = "UNSTABLE",
        global = true,
        arg_enum
    )]
    pub unstable_features: Vec<UnstableOptions>,

    /// Git repository location
    ///
    /// Without any other information, cargo will use latest commit on the main branch.
    #[clap(long, value_name = "URI", help_heading = "UNSTABLE")]
    pub git: Option<String>,

    /// Git branch to download the crate from.
    #[clap(
        long,
        value_name = "BRANCH",
        help_heading = "UNSTABLE",
        requires = "git",
        group = "git-ref"
    )]
    pub branch: Option<String>,

    /// Git tag to download the crate from.
    #[clap(
        long,
        value_name = "TAG",
        help_heading = "UNSTABLE",
        requires = "git",
        group = "git-ref"
    )]
    pub tag: Option<String>,

    /// Git reference to download the crate from
    ///
    /// This is the catch all, handling hashes to named references in remote repositories.
    #[clap(
        long,
        value_name = "REV",
        help_heading = "UNSTABLE",
        requires = "git",
        group = "git-ref"
    )]
    pub rev: Option<String>,
}

impl Args {
    /// Get dependency section
    pub fn get_section(&self) -> Vec<String> {
        if self.dev {
            vec!["dev-dependencies".to_owned()]
        } else if self.build {
            vec!["build-dependencies".to_owned()]
        } else if let Some(ref target) = self.target {
            assert!(!target.is_empty(), "Target specification may not be empty");

            vec![
                "target".to_owned(),
                target.clone(),
                "dependencies".to_owned(),
            ]
        } else {
            vec!["dependencies".to_owned()]
        }
    }

    pub fn default_features(&self) -> Option<bool> {
        resolve_bool_arg(self.default_features, self.no_default_features)
    }

    /// Build dependencies from arguments
    pub fn parse_dependencies(&self) -> Result<Vec<Dependency>> {
        let workspace_members = workspace_members(self.manifest_path.as_deref())?;

        if self.crates.len() > 1 && self.git.is_some() {
            return Err(ErrorKind::MultipleCratesWithGitOrPathOrVers.into());
        }

        if self.crates.len() > 1 && self.rename.is_some() {
            return Err(ErrorKind::MultipleCratesWithRename.into());
        }

        if self.crates.len() > 1 && self.features.is_some() {
            return Err(ErrorKind::MultipleCratesWithFeatures.into());
        }

        self.crates
            .iter()
            .map(|crate_spec| {
                self.parse_single_dependency(crate_spec, &workspace_members)
                    .map(|x| {
                        let x = self.populate_dependency(x);
                        x
                    })
            })
            .collect()
    }

    fn populate_dependency(&self, mut dependency: Dependency) -> Dependency {
        let requested_features: Option<Vec<_>> = self.features.as_ref().map(|v| {
            v.iter()
                .flat_map(|s| s.split(' '))
                .flat_map(|s| s.split(','))
                .filter(|s| !s.is_empty())
                .map(|f| f.to_owned())
                .collect()
        });

        dependency = dependency
            .set_optional(self.optional())
            .set_default_features(self.default_features())
            .set_features(requested_features);

        if let Some(ref rename) = self.rename {
            dependency = dependency.set_rename(rename);
        }

        dependency
    }

    fn parse_single_dependency(
        &self,
        crate_spec: &str,
        workspace_members: &[Package],
    ) -> Result<Dependency> {
        let crate_spec = CrateSpec::resolve(crate_spec)?;
        let manifest_path = find(&self.manifest_path)?;
        let registry_url = registry_url(&manifest_path, self.registry.as_deref())?;

        let mut dependency = match &crate_spec {
            CrateSpec::PkgId {
                name: _,
                version_req: Some(_),
            } => {
                let mut dependency = crate_spec.to_dependency()?;
                // crate specifier includes a version (e.g. `docopt@0.8`)
                if let Some(ref url) = self.git {
                    let url = url.clone();
                    let version = dependency.version().unwrap().to_string();
                    return Err(ErrorKind::GitUrlWithVersion(url, version).into());
                }

                let features = get_features_from_registry(
                    &dependency.name,
                    dependency
                        .version()
                        .expect("version populated by `parse_as_version`"),
                    &registry_url,
                )?;
                dependency = dependency.set_available_features(features);

                dependency
            }
            CrateSpec::PkgId {
                name,
                version_req: None,
            } => {
                let mut dependency = crate_spec.to_dependency()?;

                if let Some(repo) = &self.git {
                    assert!(self.registry.is_none());
                    let features = get_manifest_from_url(repo)?
                        .map(|m| m.features())
                        .transpose()?
                        .unwrap_or_else(Vec::new);

                    dependency = dependency
                        .set_git(
                            repo,
                            self.branch.clone(),
                            self.tag.clone(),
                            self.rev.clone(),
                        )
                        .set_available_features(features);
                } else if let Some(package) = workspace_members.iter().find(|p| p.name == *name) {
                    // Only special-case workspaces when the user doesn't provide any extra
                    // information, otherwise, trust the user.
                    dependency = dependency.set_path(
                        package
                            .manifest_path
                            .parent()
                            .expect("at least parent dir")
                            .as_std_path()
                            .to_owned(),
                    );
                    // dev-dependencies do not need the version populated
                    if !self.dev {
                        let op = "";
                        let v = format!("{op}{version}", op = op, version = package.version);
                        dependency = dependency.set_version(&v);
                    }
                } else {
                    let latest =
                        get_latest_dependency(name, false, &manifest_path, Some(&registry_url))?;
                    let op = "";
                    let v = format!(
                        "{op}{version}",
                        op = op,
                        // If version is unavailable `get_latest_dependency` must have
                        // returned `Err(FetchVersionError::GetVersion)`
                        version = latest.version().unwrap_or_else(|| unreachable!())
                    );
                    dependency = dependency
                        .set_version(&v)
                        .set_available_features(latest.available_features);
                }

                dependency
            }
            CrateSpec::Path(_) => {
                let mut dependency = crate_spec.to_dependency()?;
                // dev-dependencies do not need the version populated
                if !self.dev {
                    let dep_path = dependency.path().map(ToOwned::to_owned);
                    if let Some(dep_path) = dep_path {
                        if let Some(package) = workspace_members.iter().find(|p| {
                            p.manifest_path.parent().map(|p| p.as_std_path())
                                == Some(dep_path.as_path())
                        }) {
                            let op = "";
                            let v = format!("{op}{version}", op = op, version = package.version);

                            dependency = dependency.set_version(&v);
                        }
                    }
                }
                dependency
            }
        };

        if let Some(registry) = &self.registry {
            dependency = dependency.set_registry(registry);
        }

        Ok(dependency)
    }
}

impl Args {
    pub fn optional(&self) -> Option<bool> {
        resolve_bool_arg(self.optional, self.no_optional)
    }
}

#[cfg(test)]
impl Default for Args {
    fn default() -> Args {
        Args {
            crates: vec!["demo".to_owned()],
            rename: None,
            dev: false,
            build: false,
            git: None,
            branch: None,
            tag: None,
            rev: None,
            target: None,
            optional: false,
            no_optional: false,
            manifest_path: None,
            pkgid: None,
            features: None,
            no_default_features: false,
            default_features: false,
            quiet: false,
            offline: true,
            registry: None,
            unstable_features: vec![],
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, clap::ArgEnum)]
pub enum UnstableOptions {
    Git,
}

fn resolve_bool_arg(yes: bool, no: bool) -> Option<bool> {
    match (yes, no) {
        (true, false) => Some(true),
        (false, true) => Some(false),
        (false, false) => None,
        (_, _) => unreachable!("clap should make this impossible"),
    }
}

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

    #[test]
    #[cfg(feature = "test-external-apis")]
    fn test_repo_as_arg_parsing() {
        let github_url = "https://github.com/killercup/cargo-edit/";
        let args_github = Args {
            crates: vec![github_url.to_owned()],
            ..Args::default()
        };
        assert_eq!(
            args_github.parse_dependencies().unwrap(),
            vec![Dependency::new("cargo-edit").set_git(github_url, None)]
        );

        let gitlab_url = "https://gitlab.com/Polly-lang/Polly.git";
        let args_gitlab = Args {
            crates: vec![gitlab_url.to_owned()],
            ..Args::default()
        };
        assert_eq!(
            args_gitlab.parse_dependencies(None).unwrap(),
            vec![Dependency::new("polly").set_git(gitlab_url, None)]
        );
    }

    #[test]
    fn test_path_as_arg_parsing() {
        let self_path = dunce::canonicalize(std::env::current_dir().unwrap()).unwrap();
        let args_path = Args {
            // Hacky to `display` but should generally work
            crates: vec![self_path.display().to_string()],
            ..Args::default()
        };
        assert_eq!(
            args_path.parse_dependencies().unwrap()[0].path().unwrap(),
            self_path
        );
    }

    #[test]
    fn verify_app() {
        use clap::IntoApp;
        Command::into_app().debug_assert()
    }
}
<commit_msg>refactor(add): Immediately populate dependencies<commit_after>//! Handle `cargo add` arguments

#![allow(clippy::bool_assert_comparison)]

use cargo_edit::{
    find, get_features_from_registry, get_manifest_from_url, registry_url, workspace_members,
    Dependency,
};
use cargo_edit::{get_latest_dependency, CrateSpec};
use cargo_metadata::Package;
use clap::Parser;
use std::path::PathBuf;

use crate::errors::*;

#[derive(Debug, Parser)]
#[clap(bin_name = "cargo")]
pub enum Command {
    /// Add dependencies to a Cargo.toml manifest file.
    #[clap(name = "add")]
    #[clap(after_help = "\
Examples:
  $ cargo add regex
  $ cargo add regex@0.1.41 --build
  $ cargo add trycmd --dev
  $ cargo add ./crate/parser/
")]
    Add(Args),
}

#[derive(Debug, Parser)]
#[clap(about, version)]
#[clap(setting = clap::AppSettings::DeriveDisplayOrder)]
pub struct Args {
    /// Reference to a package to add as a dependency
    ///
    /// You can reference a packages by:{n}
    /// - `<name>`, like `cargo add serde` (latest version will be used){n}
    /// - `<name>@<version-req>`, like `cargo add serde@1` or `cargo add serde@=1.0.38`{n}
    /// - `<path>`, like `cargo add ./crates/parser/`
    #[clap(value_name = "DEP_ID", required = true)]
    pub crates: Vec<String>,

    /// Disable the default features
    #[clap(long)]
    no_default_features: bool,
    /// Re-enable the default features
    #[clap(long, overrides_with = "no-default-features")]
    default_features: bool,

    /// Space-separated list of features to add
    #[clap(long)]
    pub features: Option<Vec<String>>,

    /// Mark the dependency as optional
    ///
    /// The package name will be exposed as feature of your crate.
    #[clap(long, conflicts_with = "dev")]
    pub optional: bool,

    /// Mark the dependency as required
    ///
    /// The package will be removed from your features.
    #[clap(long, conflicts_with = "dev", overrides_with = "optional")]
    pub no_optional: bool,

    /// Rename the dependency
    ///
    /// Example uses:{n}
    /// - Depending on multiple versions of a crate{n}
    /// - Depend on crates with the same name from different registries
    #[clap(long, short)]
    pub rename: Option<String>,

    /// Package registry for this dependency
    #[clap(long, conflicts_with = "git")]
    pub registry: Option<String>,

    /// Add as development dependency
    ///
    /// Dev-dependencies are not used when compiling a package for building, but are used for compiling tests, examples, and benchmarks.
    ///
    /// These dependencies are not propagated to other packages which depend on this package.
    #[clap(short = 'D', long, help_heading = "SECTION", group = "section")]
    pub dev: bool,

    /// Add as build dependency
    ///
    /// Build-dependencies are the only dependencies available for use by build scripts (`build.rs`
    /// files).
    #[clap(short = 'B', long, help_heading = "SECTION", group = "section")]
    pub build: bool,

    /// Add as dependency to the given target platform.
    #[clap(
        long,
        forbid_empty_values = true,
        help_heading = "SECTION",
        group = "section"
    )]
    pub target: Option<String>,

    /// Path to `Cargo.toml`
    #[clap(long, value_name = "PATH", parse(from_os_str))]
    pub manifest_path: Option<PathBuf>,

    /// Package to modify
    #[clap(short = 'p', long = "package", value_name = "PKGID")]
    pub pkgid: Option<String>,

    /// Run without accessing the network
    #[clap(long)]
    pub offline: bool,

    /// Do not print any output in case of success.
    #[clap(long)]
    pub quiet: bool,

    /// Unstable (nightly-only) flags
    #[clap(
        short = 'Z',
        value_name = "FLAG",
        help_heading = "UNSTABLE",
        global = true,
        arg_enum
    )]
    pub unstable_features: Vec<UnstableOptions>,

    /// Git repository location
    ///
    /// Without any other information, cargo will use latest commit on the main branch.
    #[clap(long, value_name = "URI", help_heading = "UNSTABLE")]
    pub git: Option<String>,

    /// Git branch to download the crate from.
    #[clap(
        long,
        value_name = "BRANCH",
        help_heading = "UNSTABLE",
        requires = "git",
        group = "git-ref"
    )]
    pub branch: Option<String>,

    /// Git tag to download the crate from.
    #[clap(
        long,
        value_name = "TAG",
        help_heading = "UNSTABLE",
        requires = "git",
        group = "git-ref"
    )]
    pub tag: Option<String>,

    /// Git reference to download the crate from
    ///
    /// This is the catch all, handling hashes to named references in remote repositories.
    #[clap(
        long,
        value_name = "REV",
        help_heading = "UNSTABLE",
        requires = "git",
        group = "git-ref"
    )]
    pub rev: Option<String>,
}

impl Args {
    /// Get dependency section
    pub fn get_section(&self) -> Vec<String> {
        if self.dev {
            vec!["dev-dependencies".to_owned()]
        } else if self.build {
            vec!["build-dependencies".to_owned()]
        } else if let Some(ref target) = self.target {
            assert!(!target.is_empty(), "Target specification may not be empty");

            vec![
                "target".to_owned(),
                target.clone(),
                "dependencies".to_owned(),
            ]
        } else {
            vec!["dependencies".to_owned()]
        }
    }

    pub fn default_features(&self) -> Option<bool> {
        resolve_bool_arg(self.default_features, self.no_default_features)
    }

    /// Build dependencies from arguments
    pub fn parse_dependencies(&self) -> Result<Vec<Dependency>> {
        let workspace_members = workspace_members(self.manifest_path.as_deref())?;

        if self.crates.len() > 1 && self.git.is_some() {
            return Err(ErrorKind::MultipleCratesWithGitOrPathOrVers.into());
        }

        if self.crates.len() > 1 && self.rename.is_some() {
            return Err(ErrorKind::MultipleCratesWithRename.into());
        }

        if self.crates.len() > 1 && self.features.is_some() {
            return Err(ErrorKind::MultipleCratesWithFeatures.into());
        }

        self.crates
            .iter()
            .map(|crate_spec| self.parse_single_dependency(crate_spec, &workspace_members))
            .collect()
    }

    fn parse_single_dependency(
        &self,
        crate_spec: &str,
        workspace_members: &[Package],
    ) -> Result<Dependency> {
        let crate_spec = CrateSpec::resolve(crate_spec)?;
        let manifest_path = find(&self.manifest_path)?;
        let registry_url = registry_url(&manifest_path, self.registry.as_deref())?;

        let mut dependency = match &crate_spec {
            CrateSpec::PkgId {
                name: _,
                version_req: Some(_),
            } => {
                let mut dependency = crate_spec.to_dependency()?;
                dependency = self.populate_dependency(dependency);
                // crate specifier includes a version (e.g. `docopt@0.8`)
                if let Some(ref url) = self.git {
                    let url = url.clone();
                    let version = dependency.version().unwrap().to_string();
                    return Err(ErrorKind::GitUrlWithVersion(url, version).into());
                }

                let features = get_features_from_registry(
                    &dependency.name,
                    dependency
                        .version()
                        .expect("version populated by `parse_as_version`"),
                    &registry_url,
                )?;
                dependency = dependency.set_available_features(features);

                dependency
            }
            CrateSpec::PkgId {
                name,
                version_req: None,
            } => {
                let mut dependency = crate_spec.to_dependency()?;
                dependency = self.populate_dependency(dependency);

                if let Some(repo) = &self.git {
                    assert!(self.registry.is_none());
                    let features = get_manifest_from_url(repo)?
                        .map(|m| m.features())
                        .transpose()?
                        .unwrap_or_else(Vec::new);

                    dependency = dependency
                        .set_git(
                            repo,
                            self.branch.clone(),
                            self.tag.clone(),
                            self.rev.clone(),
                        )
                        .set_available_features(features);
                } else if let Some(package) = workspace_members.iter().find(|p| p.name == *name) {
                    // Only special-case workspaces when the user doesn't provide any extra
                    // information, otherwise, trust the user.
                    dependency = dependency.set_path(
                        package
                            .manifest_path
                            .parent()
                            .expect("at least parent dir")
                            .as_std_path()
                            .to_owned(),
                    );
                    // dev-dependencies do not need the version populated
                    if !self.dev {
                        let op = "";
                        let v = format!("{op}{version}", op = op, version = package.version);
                        dependency = dependency.set_version(&v);
                    }
                } else {
                    let latest =
                        get_latest_dependency(name, false, &manifest_path, Some(&registry_url))?;
                    let op = "";
                    let v = format!(
                        "{op}{version}",
                        op = op,
                        // If version is unavailable `get_latest_dependency` must have
                        // returned `Err(FetchVersionError::GetVersion)`
                        version = latest.version().unwrap_or_else(|| unreachable!())
                    );
                    dependency = dependency
                        .set_version(&v)
                        .set_available_features(latest.available_features);
                }

                dependency
            }
            CrateSpec::Path(_) => {
                let mut dependency = crate_spec.to_dependency()?;
                dependency = self.populate_dependency(dependency);

                // dev-dependencies do not need the version populated
                if !self.dev {
                    let dep_path = dependency.path().map(ToOwned::to_owned);
                    if let Some(dep_path) = dep_path {
                        if let Some(package) = workspace_members.iter().find(|p| {
                            p.manifest_path.parent().map(|p| p.as_std_path())
                                == Some(dep_path.as_path())
                        }) {
                            let op = "";
                            let v = format!("{op}{version}", op = op, version = package.version);

                            dependency = dependency.set_version(&v);
                        }
                    }
                }
                dependency
            }
        };

        if let Some(registry) = &self.registry {
            dependency = dependency.set_registry(registry);
        }

        Ok(dependency)
    }

    fn populate_dependency(&self, mut dependency: Dependency) -> Dependency {
        let requested_features: Option<Vec<_>> = self.features.as_ref().map(|v| {
            v.iter()
                .flat_map(|s| s.split(' '))
                .flat_map(|s| s.split(','))
                .filter(|s| !s.is_empty())
                .map(|f| f.to_owned())
                .collect()
        });

        dependency = dependency
            .set_optional(self.optional())
            .set_default_features(self.default_features())
            .set_features(requested_features);

        if let Some(ref rename) = self.rename {
            dependency = dependency.set_rename(rename);
        }

        dependency
    }
}

impl Args {
    pub fn optional(&self) -> Option<bool> {
        resolve_bool_arg(self.optional, self.no_optional)
    }
}

#[cfg(test)]
impl Default for Args {
    fn default() -> Args {
        Args {
            crates: vec!["demo".to_owned()],
            rename: None,
            dev: false,
            build: false,
            git: None,
            branch: None,
            tag: None,
            rev: None,
            target: None,
            optional: false,
            no_optional: false,
            manifest_path: None,
            pkgid: None,
            features: None,
            no_default_features: false,
            default_features: false,
            quiet: false,
            offline: true,
            registry: None,
            unstable_features: vec![],
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, clap::ArgEnum)]
pub enum UnstableOptions {
    Git,
}

fn resolve_bool_arg(yes: bool, no: bool) -> Option<bool> {
    match (yes, no) {
        (true, false) => Some(true),
        (false, true) => Some(false),
        (false, false) => None,
        (_, _) => unreachable!("clap should make this impossible"),
    }
}

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

    #[test]
    #[cfg(feature = "test-external-apis")]
    fn test_repo_as_arg_parsing() {
        let github_url = "https://github.com/killercup/cargo-edit/";
        let args_github = Args {
            crates: vec![github_url.to_owned()],
            ..Args::default()
        };
        assert_eq!(
            args_github.parse_dependencies().unwrap(),
            vec![Dependency::new("cargo-edit").set_git(github_url, None)]
        );

        let gitlab_url = "https://gitlab.com/Polly-lang/Polly.git";
        let args_gitlab = Args {
            crates: vec![gitlab_url.to_owned()],
            ..Args::default()
        };
        assert_eq!(
            args_gitlab.parse_dependencies(None).unwrap(),
            vec![Dependency::new("polly").set_git(gitlab_url, None)]
        );
    }

    #[test]
    fn test_path_as_arg_parsing() {
        let self_path = dunce::canonicalize(std::env::current_dir().unwrap()).unwrap();
        let args_path = Args {
            // Hacky to `display` but should generally work
            crates: vec![self_path.display().to_string()],
            ..Args::default()
        };
        assert_eq!(
            args_path.parse_dependencies().unwrap()[0].path().unwrap(),
            self_path
        );
    }

    #[test]
    fn verify_app() {
        use clap::IntoApp;
        Command::into_app().debug_assert()
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(test, allow(dead_code))]

#![feature(collections)]
#![feature(env)]
#![feature(old_io)]
#![feature(old_path)]
#![feature(rustc_private)]

extern crate document;
extern crate xpath;
extern crate getopts;

use std::cmp::min;
use std::borrow::ToOwned;
use std::env;
use std::collections::HashMap;
use std::old_io::File;
use std::old_io::stdio;

use document::parser::Parser;

use xpath::nodeset::ToNode;
use xpath::{EvaluationContext,Factory,Expression,Value};

use getopts::{reqopt,optmulti,getopts,OptGroup,usage};

fn print_usage(program: &str, opts: &[OptGroup]) {
    let brief = format!("Usage: {} [options] FILE...", program);
    print!("{}", usage(&brief, opts));
}

fn pretty_error(xml: &str, position: usize) -> &str {
    let s = &xml[position..];
    let l = s.chars().count();
    s.slice_chars(0, min(l, 15))
}

fn build_xpath(xpath_str: &str) -> Box<Expression> {
    let factory = Factory::new();

    match factory.build(xpath_str) {
        Err(x) => panic!("Unable to compile XPath: {:?}", x),
        Ok(None) => panic!("Unable to compile XPath"),
        Ok(Some(x)) => x,
    }
}

fn load_xml<R>(input: R) -> document::Package
    where R: Reader
{
    let mut input = input;
    let p = Parser::new();

    let data = match input.read_to_string() {
        Ok(x) => x,
        Err(x) => panic!("Can't read: {}", x),
    };

    match p.parse(&data) {
        Ok(d) => d,
        Err(point) => panic!("Unable to parse: {}", pretty_error(&data, point)),
    }
}

fn argument_name_value(s: &str, delim: char) -> Option<(&str, &str)> {
    let parts: Vec<_> = s.splitn(2, delim).collect();
    if parts.len() < 2 {
        None
    } else {
        Some((parts[0], parts[1]))
    }
}

fn build_functions() -> xpath::Functions {
    let mut functions = HashMap::new();
    xpath::function::register_core_functions(&mut functions);
    functions
}

fn build_variables<'a>(arguments: &getopts::Matches) -> xpath::Variables<'a> {
    let mut vars = HashMap::new();

    for boolean in arguments.opt_strs("boolean") {
        match argument_name_value(&boolean, '=') {
            Some((name, "true")) => { vars.insert(name.to_owned(), Value::Boolean(true)); },
            Some((name, "false")) => { vars.insert(name.to_owned(), Value::Boolean(false)); },
            Some((_, val)) => panic!("Unknown boolean value '{}'", val),
            None => panic!("boolean argument '{}' is malformed", boolean),
        }
    }

    for number in arguments.opt_strs("number") {
        match argument_name_value(&number, '=') {
            Some((name, val)) => {
                match val.parse() {
                    Ok(val) => { vars.insert(name.to_owned(), Value::Number(val)); },
                    Err(e) => panic!("Unknown numeric value '{}': {}", val, e),
                }
            },
            None => panic!("number argument '{}' is malformed", number),
        }
    }

    for string in arguments.opt_strs("string") {
        match argument_name_value(&string, '=') {
            Some((name, val)) => { vars.insert(name.to_owned(), Value::String(val.to_owned())); },
            None => panic!("string argument '{}' is malformed", string),
        }
    }

    vars
}

fn build_namespaces(arguments: &getopts::Matches) -> xpath::Namespaces {
    let mut namespaces = HashMap::new();

    for ns_defn in arguments.opt_strs("namespace") {
        match argument_name_value(&ns_defn, ':') {
            Some((prefix, uri)) => { namespaces.insert(prefix.to_owned(), uri.to_owned()); },
            None => panic!("Namespace definition '{}' is malformed", ns_defn),
        }
    }

    namespaces
}

fn main() {
    let args: Vec<_> = env::args().collect();

    let program_name = &args[0];

    let opts = &[
        reqopt("", "xpath", "The XPath to execute", "XPATH"),
        optmulti("", "namespace", "set namespace prefix", "PREFIX:URI"),
        optmulti("", "string", "set string variable", "NAME=VALUE"),
        optmulti("", "number", "set number variable", "NAME=VALUE"),
        optmulti("", "boolean", "set boolean variable", "NAME=VALUE"),
    ];

    let arguments = match getopts(args.tail(), opts) {
        Ok(x) => x,
        Err(x) => {
            println!("{}", x);
            print_usage(program_name, opts);
            return;
        },
    };

    let xpath_str = arguments.opt_str("xpath").unwrap();
    let xpath = build_xpath(&xpath_str);

    for filename in &arguments.free {
        let package = if *filename == "-" {
            load_xml(stdio::stdin_raw())
        } else {
            load_xml(File::open(&Path::new(filename)))
        };

        let doc = package.as_document();
        let root = doc.root().into_node();

        let functions = build_functions();
        let variables = build_variables(&arguments);
        let namespaces = build_namespaces(&arguments);

        let context = EvaluationContext::new(
            root,
            &functions,
            &variables,
            &namespaces,
        );

        let res = xpath.evaluate(&context);

        println!("{:?}", res);
    }
}
<commit_msg>env feature is no longer gated<commit_after>#![cfg_attr(test, allow(dead_code))]

#![feature(collections)]
#![feature(old_io)]
#![feature(old_path)]
#![feature(rustc_private)]

extern crate document;
extern crate xpath;
extern crate getopts;

use std::cmp::min;
use std::borrow::ToOwned;
use std::env;
use std::collections::HashMap;
use std::old_io::File;
use std::old_io::stdio;

use document::parser::Parser;

use xpath::nodeset::ToNode;
use xpath::{EvaluationContext,Factory,Expression,Value};

use getopts::{reqopt,optmulti,getopts,OptGroup,usage};

fn print_usage(program: &str, opts: &[OptGroup]) {
    let brief = format!("Usage: {} [options] FILE...", program);
    print!("{}", usage(&brief, opts));
}

fn pretty_error(xml: &str, position: usize) -> &str {
    let s = &xml[position..];
    let l = s.chars().count();
    s.slice_chars(0, min(l, 15))
}

fn build_xpath(xpath_str: &str) -> Box<Expression> {
    let factory = Factory::new();

    match factory.build(xpath_str) {
        Err(x) => panic!("Unable to compile XPath: {:?}", x),
        Ok(None) => panic!("Unable to compile XPath"),
        Ok(Some(x)) => x,
    }
}

fn load_xml<R>(input: R) -> document::Package
    where R: Reader
{
    let mut input = input;
    let p = Parser::new();

    let data = match input.read_to_string() {
        Ok(x) => x,
        Err(x) => panic!("Can't read: {}", x),
    };

    match p.parse(&data) {
        Ok(d) => d,
        Err(point) => panic!("Unable to parse: {}", pretty_error(&data, point)),
    }
}

fn argument_name_value(s: &str, delim: char) -> Option<(&str, &str)> {
    let parts: Vec<_> = s.splitn(2, delim).collect();
    if parts.len() < 2 {
        None
    } else {
        Some((parts[0], parts[1]))
    }
}

fn build_functions() -> xpath::Functions {
    let mut functions = HashMap::new();
    xpath::function::register_core_functions(&mut functions);
    functions
}

fn build_variables<'a>(arguments: &getopts::Matches) -> xpath::Variables<'a> {
    let mut vars = HashMap::new();

    for boolean in arguments.opt_strs("boolean") {
        match argument_name_value(&boolean, '=') {
            Some((name, "true")) => { vars.insert(name.to_owned(), Value::Boolean(true)); },
            Some((name, "false")) => { vars.insert(name.to_owned(), Value::Boolean(false)); },
            Some((_, val)) => panic!("Unknown boolean value '{}'", val),
            None => panic!("boolean argument '{}' is malformed", boolean),
        }
    }

    for number in arguments.opt_strs("number") {
        match argument_name_value(&number, '=') {
            Some((name, val)) => {
                match val.parse() {
                    Ok(val) => { vars.insert(name.to_owned(), Value::Number(val)); },
                    Err(e) => panic!("Unknown numeric value '{}': {}", val, e),
                }
            },
            None => panic!("number argument '{}' is malformed", number),
        }
    }

    for string in arguments.opt_strs("string") {
        match argument_name_value(&string, '=') {
            Some((name, val)) => { vars.insert(name.to_owned(), Value::String(val.to_owned())); },
            None => panic!("string argument '{}' is malformed", string),
        }
    }

    vars
}

fn build_namespaces(arguments: &getopts::Matches) -> xpath::Namespaces {
    let mut namespaces = HashMap::new();

    for ns_defn in arguments.opt_strs("namespace") {
        match argument_name_value(&ns_defn, ':') {
            Some((prefix, uri)) => { namespaces.insert(prefix.to_owned(), uri.to_owned()); },
            None => panic!("Namespace definition '{}' is malformed", ns_defn),
        }
    }

    namespaces
}

fn main() {
    let args: Vec<_> = env::args().collect();

    let program_name = &args[0];

    let opts = &[
        reqopt("", "xpath", "The XPath to execute", "XPATH"),
        optmulti("", "namespace", "set namespace prefix", "PREFIX:URI"),
        optmulti("", "string", "set string variable", "NAME=VALUE"),
        optmulti("", "number", "set number variable", "NAME=VALUE"),
        optmulti("", "boolean", "set boolean variable", "NAME=VALUE"),
    ];

    let arguments = match getopts(args.tail(), opts) {
        Ok(x) => x,
        Err(x) => {
            println!("{}", x);
            print_usage(program_name, opts);
            return;
        },
    };

    let xpath_str = arguments.opt_str("xpath").unwrap();
    let xpath = build_xpath(&xpath_str);

    for filename in &arguments.free {
        let package = if *filename == "-" {
            load_xml(stdio::stdin_raw())
        } else {
            load_xml(File::open(&Path::new(filename)))
        };

        let doc = package.as_document();
        let root = doc.root().into_node();

        let functions = build_functions();
        let variables = build_variables(&arguments);
        let namespaces = build_namespaces(&arguments);

        let context = EvaluationContext::new(
            root,
            &functions,
            &variables,
            &namespaces,
        );

        let res = xpath.evaluate(&context);

        println!("{:?}", res);
    }
}
<|endoftext|>


<commit_before>const NIL: char = '';

/// The mapping for [code page 850](https://en.wikipedia.org/wiki/Code_page_850)
///
/// This code page is also known as `DOS/IBM-ASCII` and is used as the default
/// code page by this library.
pub const CODEPAGE_0850: &'static [char]  = &[
//   0     1     2     3     4     5     6     7     8     9     A     B     C     D     E     F
    NIL,  '☺',  '☻',  '♥',  '♦',  '♣',  '♠',  '•',  '◘',  '○',  '◙',  '♂',  '♀',  '♪',  '♫',  '☼', // 0
    '►',  '◄',  '↕',  '‼',  '¶',  '§',  '▬',  '↨',  '↑',  '↓',  '→',  '←',  '∟',  '↔',  '▲',  '▼', // 1
    NIL,  '!',  '"',  '#',  '$',  '%',  '&', '\'',  '(',  ')',  '*',  '+',  ',',  '-',  '.',  '/', // 2
    '0',  '1',  '2',  '3',  '4',  '5',  '6',  '7',  '8',  '9',  ':',  ';',  '<',  '=',  '>',  '?', // 3
    '@',  'A',  'B',  'C',  'D',  'E',  'F',  'G',  'H',  'I',  'J',  'K',  'L',  'M',  'N',  'O', // 4
    'P',  'Q',  'R',  'S',  'T',  'U',  'V',  'W',  'X',  'Y',  'Z',  '[', '\\',  ']',  '^',  '_', // 5
    '`',  'a',  'b',  'c',  'd',  'e',  'f',  'g',  'h',  'i',  'j',  'k',  'l',  'm',  'n',  'o', // 6
    'p',  'q',  'r',  's',  't',  'u',  'v',  'w',  'x',  'y',  'z',  '{',  '|',  '}',  '~',  '⌂', // 7
    'Ç',  'ü',  'é',  'â',  'ä',  'à',  'å',  'ç',  'ê',  'ë',  'è',  'ï',  'î',  'ì',  'Ä',  'Å', // 8
    'É',  'æ',  'Æ',  'ô',  'ö',  'ò',  'û',  'ù',  'ÿ',  'Ö',  'Ü',  'ø',  '£',  'Ø',  '×',  'ƒ', // 9
    'á',  'í',  'ó',  'ú',  'ñ',  'Ñ',  'ª',  'º',  '¿',  '⌐',  '¬',  '½',  '¼',  '¡',  '«',  '»', // A
    '░',  '▒',  '▓',  '│',  '┤',  '╡',  '╢',  '╖',  '╕',  '╣',  '║',  '╗',  '╝',  '╜',  '╛',  '┐', // B
    '└',  '┴',  '┬',  '├',  '─',  '┼',  '╞',  '╟',  '╚',  '╔',  '╩',  '╦',  '╠',  '═',  '╬',  '╧', // C
    '╨',  '╤',  '╥',  '╙',  '╘',  '╒',  '╕',  '╫',  '╪',  '┘',  '┌',  '█',  '▄',  '▌',  '▐',  '▀', // D
    'α',  'ß',  'Γ',  '∏',  '∑',  'σ',  'µ',  'τ',  'Φ',  'θ',  'Ω',  'δ',  '∞',  'Ø',  'ε',  '∩', // E
    '≡',  '±',  '≥',  '≤',  '⌠',  '⌡',  '÷',  '≈',  '▫',  '¨',  '·',  '√',  'ⁿ',  '²',  '■',  NIL, // F
];

/// The mapping for [code page 1252](https://en.wikipedia.org/wiki/Code_page_1252)
///
/// This code page is also known as `Latin 1 Windows` or `ANSI`.
pub const CODEPAGE_1252: &'static [char] = &[
//   0     1     2     3     4     5     6     7     8     9     A     B     C     D     E     F
    NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL, // 0
    NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL, // 1
    ' ',  '!',  '"',  '#',  '$',  '%',  '&', '\'',  '(',  ')',  '*',  '+',  ',',  '-',  '.',  '/', // 2
    '0',  '1',  '2',  '3',  '4',  '5',  '6',  '7',  '8',  '9',  ':',  ';',  '<',  '=',  '>',  '?', // 3
    '@',  'A',  'B',  'C',  'D',  'E',  'F',  'G',  'H',  'I',  'J',  'K',  'L',  'M',  'N',  'O', // 4
    'P',  'Q',  'R',  'S',  'T',  'U',  'V',  'W',  'X',  'Y',  'Z',  '[', '\\',  ']',  '^',  '_', // 5
    '`',  'a',  'b',  'c',  'd',  'e',  'f',  'g',  'h',  'i',  'j',  'k',  'l',  'm',  'n',  'o', // 6
    'p',  'q',  'r',  's',  't',  'u',  'v',  'w',  'x',  'y',  'z',  '{',  '|',  '}',  '~',  NIL, // 7
    '€',  NIL,  '‚',  'ƒ',  '„',  '…',  '†',  '‡',  'ˆ',  '‰',  'Š',  '‹',  'Œ',  NIL,  'Ž',  NIL, // 8
    NIL,  '‘',  '’',  '“',  '”',  '•',  '–',  '—',  '˜',  '™',  'Š',  '›',  'œ',  NIL,  'ž',  'Ÿ', // 9
    ' ',  '¡',  '¢',  '£',  '¤',  '¥',  '¦',  '§',  '¨',  '©',  'ª',  '«',  '¬',  NIL,  '®',  '¯', // A
    '°',  '±',  '²',  '³',  '´',  'µ',  '¶',  '·',  '¸',  '¹',  'º',  '»',  '¼',  '½',  '¾',  '¿', // B
    'À',  'Á',  'Â',  'Ã',  'Ä',  'Å',  'Æ',  'Ç',  'È',  'É',  'Ê',  'Ë',  'Ì',  'Í',  'Î',  'Ï', // C
    'Ð',  'Ñ',  'Ò',  'Ó',  'Ô',  'Õ',  'Ö',  '×',  'Ø',  'Ù',  'Ú',  'Û',  'Ü',  'Ý',  'Þ',  'ß', // D
    'à',  'á',  'â',  'ã',  'ä',  'å',  'æ',  'ç',  'è',  'é',  'ê',  'ë',  'ì',  'í',  'î',  'ï', // E
    'ð',  'ñ',  'ò',  'ó',  'ô',  'õ',  'ö',  '÷',  'ø',  'ù',  'ú',  'û',  'ü',  'ý',  'þ',  'ÿ', // F
];

fn contains(byte: u8, codepage: &[char]) -> bool {
    (byte as usize) < codepage.len()
}

/// Returns a byte's character representation given a specific codepage
pub fn as_char(byte: u8, codepage: &[char]) -> char {
    if !contains(byte, codepage) {
        return NIL;
    }

    codepage[byte as usize]
}
<commit_msg>[fix] Fix 0x20 byte mapping for code page 850<commit_after>const NIL: char = '';

/// The mapping for [code page 850](https://en.wikipedia.org/wiki/Code_page_850)
///
/// This code page is also known as `DOS/IBM-ASCII` and is used as the default
/// code page by this library.
pub const CODEPAGE_0850: &'static [char]  = &[
//   0     1     2     3     4     5     6     7     8     9     A     B     C     D     E     F
    NIL,  '☺',  '☻',  '♥',  '♦',  '♣',  '♠',  '•',  '◘',  '○',  '◙',  '♂',  '♀',  '♪',  '♫',  '☼', // 0
    '►',  '◄',  '↕',  '‼',  '¶',  '§',  '▬',  '↨',  '↑',  '↓',  '→',  '←',  '∟',  '↔',  '▲',  '▼', // 1
    ' ',  '!',  '"',  '#',  '$',  '%',  '&', '\'',  '(',  ')',  '*',  '+',  ',',  '-',  '.',  '/', // 2
    '0',  '1',  '2',  '3',  '4',  '5',  '6',  '7',  '8',  '9',  ':',  ';',  '<',  '=',  '>',  '?', // 3
    '@',  'A',  'B',  'C',  'D',  'E',  'F',  'G',  'H',  'I',  'J',  'K',  'L',  'M',  'N',  'O', // 4
    'P',  'Q',  'R',  'S',  'T',  'U',  'V',  'W',  'X',  'Y',  'Z',  '[', '\\',  ']',  '^',  '_', // 5
    '`',  'a',  'b',  'c',  'd',  'e',  'f',  'g',  'h',  'i',  'j',  'k',  'l',  'm',  'n',  'o', // 6
    'p',  'q',  'r',  's',  't',  'u',  'v',  'w',  'x',  'y',  'z',  '{',  '|',  '}',  '~',  '⌂', // 7
    'Ç',  'ü',  'é',  'â',  'ä',  'à',  'å',  'ç',  'ê',  'ë',  'è',  'ï',  'î',  'ì',  'Ä',  'Å', // 8
    'É',  'æ',  'Æ',  'ô',  'ö',  'ò',  'û',  'ù',  'ÿ',  'Ö',  'Ü',  'ø',  '£',  'Ø',  '×',  'ƒ', // 9
    'á',  'í',  'ó',  'ú',  'ñ',  'Ñ',  'ª',  'º',  '¿',  '⌐',  '¬',  '½',  '¼',  '¡',  '«',  '»', // A
    '░',  '▒',  '▓',  '│',  '┤',  '╡',  '╢',  '╖',  '╕',  '╣',  '║',  '╗',  '╝',  '╜',  '╛',  '┐', // B
    '└',  '┴',  '┬',  '├',  '─',  '┼',  '╞',  '╟',  '╚',  '╔',  '╩',  '╦',  '╠',  '═',  '╬',  '╧', // C
    '╨',  '╤',  '╥',  '╙',  '╘',  '╒',  '╕',  '╫',  '╪',  '┘',  '┌',  '█',  '▄',  '▌',  '▐',  '▀', // D
    'α',  'ß',  'Γ',  '∏',  '∑',  'σ',  'µ',  'τ',  'Φ',  'θ',  'Ω',  'δ',  '∞',  'Ø',  'ε',  '∩', // E
    '≡',  '±',  '≥',  '≤',  '⌠',  '⌡',  '÷',  '≈',  '▫',  '¨',  '·',  '√',  'ⁿ',  '²',  '■',  NIL, // F
];

/// The mapping for [code page 1252](https://en.wikipedia.org/wiki/Code_page_1252)
///
/// This code page is also known as `Latin 1 Windows` or `ANSI`.
pub const CODEPAGE_1252: &'static [char] = &[
//   0     1     2     3     4     5     6     7     8     9     A     B     C     D     E     F
    NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL, // 0
    NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL,  NIL, // 1
    ' ',  '!',  '"',  '#',  '$',  '%',  '&', '\'',  '(',  ')',  '*',  '+',  ',',  '-',  '.',  '/', // 2
    '0',  '1',  '2',  '3',  '4',  '5',  '6',  '7',  '8',  '9',  ':',  ';',  '<',  '=',  '>',  '?', // 3
    '@',  'A',  'B',  'C',  'D',  'E',  'F',  'G',  'H',  'I',  'J',  'K',  'L',  'M',  'N',  'O', // 4
    'P',  'Q',  'R',  'S',  'T',  'U',  'V',  'W',  'X',  'Y',  'Z',  '[', '\\',  ']',  '^',  '_', // 5
    '`',  'a',  'b',  'c',  'd',  'e',  'f',  'g',  'h',  'i',  'j',  'k',  'l',  'm',  'n',  'o', // 6
    'p',  'q',  'r',  's',  't',  'u',  'v',  'w',  'x',  'y',  'z',  '{',  '|',  '}',  '~',  NIL, // 7
    '€',  NIL,  '‚',  'ƒ',  '„',  '…',  '†',  '‡',  'ˆ',  '‰',  'Š',  '‹',  'Œ',  NIL,  'Ž',  NIL, // 8
    NIL,  '‘',  '’',  '“',  '”',  '•',  '–',  '—',  '˜',  '™',  'Š',  '›',  'œ',  NIL,  'ž',  'Ÿ', // 9
    ' ',  '¡',  '¢',  '£',  '¤',  '¥',  '¦',  '§',  '¨',  '©',  'ª',  '«',  '¬',  NIL,  '®',  '¯', // A
    '°',  '±',  '²',  '³',  '´',  'µ',  '¶',  '·',  '¸',  '¹',  'º',  '»',  '¼',  '½',  '¾',  '¿', // B
    'À',  'Á',  'Â',  'Ã',  'Ä',  'Å',  'Æ',  'Ç',  'È',  'É',  'Ê',  'Ë',  'Ì',  'Í',  'Î',  'Ï', // C
    'Ð',  'Ñ',  'Ò',  'Ó',  'Ô',  'Õ',  'Ö',  '×',  'Ø',  'Ù',  'Ú',  'Û',  'Ü',  'Ý',  'Þ',  'ß', // D
    'à',  'á',  'â',  'ã',  'ä',  'å',  'æ',  'ç',  'è',  'é',  'ê',  'ë',  'ì',  'í',  'î',  'ï', // E
    'ð',  'ñ',  'ò',  'ó',  'ô',  'õ',  'ö',  '÷',  'ø',  'ù',  'ú',  'û',  'ü',  'ý',  'þ',  'ÿ', // F
];

fn contains(byte: u8, codepage: &[char]) -> bool {
    (byte as usize) < codepage.len()
}

/// Returns a byte's character representation given a specific codepage
pub fn as_char(byte: u8, codepage: &[char]) -> char {
    if !contains(byte, codepage) {
        return NIL;
    }

    codepage[byte as usize]
}
<|endoftext|>


<commit_before>// This is a test which attempts to blow out the system limit with how many
// arguments can be passed to a process. This'll successively call rustc with
// larger and larger argument lists in an attempt to find one that's way too
// big for the system at hand. This file itself is then used as a "linker" to
// detect when the process creation succeeds.
//
// Eventually we should see an argument that looks like `@` as we switch from
// passing literal arguments to passing everything in the file.

use std::env;
use std::fs::{self, File};
use std::io::{BufWriter, Write, Read};
use std::path::PathBuf;
use std::process::Command;

fn main() {
    let tmpdir = PathBuf::from(env::var_os("TMPDIR").unwrap());
    let ok = tmpdir.join("ok");
    if env::var("YOU_ARE_A_LINKER").is_ok() {
        if let Some(file) = env::args_os().find(|a| a.to_string_lossy().contains("@")) {
            let file = file.to_str().expect("non-utf8 file argument");
            fs::copy(&file[1..], &ok).unwrap();
        }
        return
    }

    let rustc = env::var_os("RUSTC").unwrap_or("rustc".into());
    let me_as_linker = format!("linker={}", env::current_exe().unwrap().display());
    for i in (1..).map(|i| i * 100) {
        println!("attempt: {}", i);
        let file = tmpdir.join("bar.rs");
        let mut f = BufWriter::new(File::create(&file).unwrap());
        let mut lib_name = String::new();
        for _ in 0..i {
            lib_name.push_str("foo");
        }
        for j in 0..i {
            writeln!(f, "#[link(name = \"{}{}\")]", lib_name, j).unwrap();
        }
        writeln!(f, "extern {{}}\nfn main() {{}}").unwrap();
        f.into_inner().unwrap();

        drop(fs::remove_file(&ok));
        let output = Command::new(&rustc)
            .arg(&file)
            .arg("-C").arg(&me_as_linker)
            .arg("--out-dir").arg(&tmpdir)
            .env("YOU_ARE_A_LINKER", "1")
            .output()
            .unwrap();

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            panic!("status: {}\nstdout:\n{}\nstderr:\n{}",
                   output.status,
                   String::from_utf8_lossy(&output.stdout),
                   stderr.lines().map(|l| {
                       if l.len() > 200 {
                           format!("{}...\n", &l[..200])
                       } else {
                           format!("{}\n", l)
                       }
                   }).collect::<String>());
        }

        if !ok.exists() {
            continue
        }

        let mut contents = Vec::new();
        File::open(&ok).unwrap().read_to_end(&mut contents).unwrap();

        for j in 0..i {
            let exp = format!("{}{}", lib_name, j);
            let exp = if cfg!(target_env = "msvc") {
                let mut out = Vec::with_capacity(exp.len() * 2);
                for c in exp.encode_utf16() {
                    // encode in little endian
                    out.push(c as u8);
                    out.push((c >> 8) as u8);
                }
                out
            } else {
                exp.into_bytes()
            };
            assert!(contents.windows(exp.len()).any(|w| w == &exp[..]));
        }

        break
    }
}
<commit_msg>Optimize long-linker-command-line test<commit_after>// This is a test which attempts to blow out the system limit with how many
// arguments can be passed to a process. This'll successively call rustc with
// larger and larger argument lists in an attempt to find one that's way too
// big for the system at hand. This file itself is then used as a "linker" to
// detect when the process creation succeeds.
//
// Eventually we should see an argument that looks like `@` as we switch from
// passing literal arguments to passing everything in the file.

use std::collections::HashSet;
use std::env;
use std::fs::{self, File};
use std::io::{BufWriter, Write};
use std::path::{Path, PathBuf};
use std::process::Command;

fn write_test_case(file: &Path, n: usize) -> HashSet<String> {
    let mut libs = HashSet::new();
    let mut f = BufWriter::new(File::create(&file).unwrap());
    let mut prefix = String::new();
    for _ in 0..n {
        prefix.push_str("foo");
    }
    for i in 0..n {
        writeln!(f, "#[link(name = \"S{}{}S\")]", prefix, i).unwrap();
        libs.insert(format!("{}{}", prefix, i));
    }
    writeln!(f, "extern {{}}\nfn main() {{}}").unwrap();
    f.into_inner().unwrap();

    libs
}

fn read_linker_args(path: &Path) -> String {
    let contents = fs::read(path).unwrap();
    if cfg!(target_env = "msvc") {
        let mut i = contents.chunks(2).map(|c| {
            c[0] as u16 | ((c[1] as u16) << 8)
        });
        assert_eq!(i.next(), Some(0xfeff), "Expected UTF-16 BOM");
        String::from_utf16(&i.collect::<Vec<u16>>()).unwrap()
    } else {
        String::from_utf8(contents).unwrap()
    }
}

fn main() {
    let tmpdir = PathBuf::from(env::var_os("TMPDIR").unwrap());
    let ok = tmpdir.join("ok");
    if env::var("YOU_ARE_A_LINKER").is_ok() {
        if let Some(file) = env::args_os().find(|a| a.to_string_lossy().contains("@")) {
            let file = file.to_str().expect("non-utf8 file argument");
            fs::copy(&file[1..], &ok).unwrap();
        }
        return
    }

    let rustc = env::var_os("RUSTC").unwrap_or("rustc".into());
    let me_as_linker = format!("linker={}", env::current_exe().unwrap().display());
    for i in (1..).map(|i| i * 100) {
        println!("attempt: {}", i);
        let file = tmpdir.join("bar.rs");
        let mut expected_libs = write_test_case(&file, i);

        drop(fs::remove_file(&ok));
        let output = Command::new(&rustc)
            .arg(&file)
            .arg("-C").arg(&me_as_linker)
            .arg("--out-dir").arg(&tmpdir)
            .env("YOU_ARE_A_LINKER", "1")
            .output()
            .unwrap();

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            panic!("status: {}\nstdout:\n{}\nstderr:\n{}",
                   output.status,
                   String::from_utf8_lossy(&output.stdout),
                   stderr.lines().map(|l| {
                       if l.len() > 200 {
                           format!("{}...\n", &l[..200])
                       } else {
                           format!("{}\n", l)
                       }
                   }).collect::<String>());
        }

        if !ok.exists() {
            continue
        }

        let linker_args = read_linker_args(&ok);
        for mut arg in linker_args.split('S') {
            expected_libs.remove(arg);
        }

        assert!(
            expected_libs.is_empty(),
            "expected but missing libraries: {:#?}\nlinker arguments: \n{}",
            expected_libs,
            linker_args,
        );

        break
    }
}
<|endoftext|>


<commit_before>use data::packet::Packet;
use data::rational::Rational64;
use stream::Stream;
use error::*;

use buffer::Buffered;
use std::io::SeekFrom;

#[derive(Clone, Debug, PartialEq)]
pub struct GlobalInfo {
    pub duration: Option<u64>,
    pub timebase: Option<Rational64>,
    pub streams: Vec<Stream>,
}

#[derive(Clone, Debug, PartialEq)]
pub enum Event {
    NewPacket(Packet),
    NewStream(Stream),
    MoreDataNeeded(usize),
}

pub trait Demuxer : Send {
    fn read_headers(&mut self, buf: &Box<Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom>;
    fn read_event(&mut self, buf: &Box<Buffered>) -> Result<(SeekFrom, Event)>;
}

#[derive(Clone, Debug, PartialEq)]
pub struct Descr {
    pub name: &'static str,
    pub demuxer: &'static str,
    pub description: &'static str,
    pub extensions: &'static [&'static str],
    pub mime: &'static [&'static str],
}

pub trait Descriptor {
    fn create(&self) -> Box<Demuxer>;
    fn describe<'a>(&'a self) -> &'a Descr;

    fn probe(&self, data: &[u8]) -> u8;
}

pub struct Context {
    demuxer: Box<Demuxer>,
    reader: Box<Buffered>,
    pub info: GlobalInfo,
}

impl Context {
    pub fn new<R: Buffered + 'static>(demuxer: Box<Demuxer>, reader: Box<R>) -> Self {
        Context {
            demuxer: demuxer,
            reader: reader,
            info: GlobalInfo {
                duration: None,
                timebase: None,
                streams: Vec::with_capacity(2),
            },
        }
    }

    fn read_headers_internal(&mut self) -> Result<()> {
        let ref mut demux = self.demuxer;
        let res = demux.read_headers(&self.reader, &mut self.info);
        match res {
            Err(e) => Err(e),
            Ok(seek) => {
                //TODO: handle seeking here
                let res = self.reader.seek(seek);
                println!("stream now at index: {:?}", res);
                Ok(())
            }
        }
    }

    pub fn read_headers(&mut self) -> Result<()> {
        loop {
            println!("Filling");
            try!(self.reader.fill_buf());
            match self.read_headers_internal() {
                Err(e) => match e {
                    Error(ErrorKind::MoreDataNeeded(needed), _) => {
                        self.reader.grow(needed);
                    },
                    _ => return Err(e)
                },
                Ok(_) => return Ok(()),
            }
        }
    }

    fn read_event_internal(&mut self) -> Result<Event> {
        let ref mut demux = self.demuxer;

        let res = demux.read_event(&self.reader);
        match res {
            Err(e) => Err(e),
            Ok((seek, event)) => {
                //TODO: handle seeking here
                let res = self.reader.seek(seek);
                if let Event::NewStream(ref st) = event {
                    self.info.streams.push(st.clone());
                }
                if let Event::MoreDataNeeded(size) = event {
                    return Err(ErrorKind::MoreDataNeeded(size).into());
                }
                Ok(event)
            }
        }
    }

    pub fn read_event(&mut self) -> Result<Event> {
        // TODO: guard against infiniloops and maybe factor the loop.
        loop {
            match self.read_event_internal() {
                Err(e) => match e {
                    Error(ErrorKind::MoreDataNeeded(needed), _) => {
                        self.reader.grow(needed);
                        try!(self.reader.fill_buf());
                    },
                    _ => return Err(e)
                },
                Ok(ev) => return Ok(ev),
            }
        }
    }
}

pub const PROBE_DATA: usize = 4 * 1024;
pub const PROBE_SCORE_EXTENSION: u8 = 50;

// TODO:
// IntoIterator<Item = &'static Descriptor> is confusing

pub trait Probe {
    fn probe(&self, data: &[u8]) -> Option<&'static Descriptor>;
}

impl<'a> Probe for [&'static Descriptor] {
    fn probe(&self, data: &[u8]) -> Option<&'static Descriptor> {
        let mut max = u8::min_value();
        let mut candidate: Option<&'static Descriptor> = None;
        for desc in self {
            let score = desc.probe(data);

            if score > max {
                max = score;
                candidate = Some(*desc);
            }
        }

        if max > PROBE_SCORE_EXTENSION {
            candidate
        } else {
            None
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::io::SeekFrom;
    use data::packet::Packet;

    struct DummyDes {
        d: Descr,
    }

    struct DummyDemuxer {}

    impl Demuxer for DummyDemuxer {
        fn read_headers(&mut self, buf: &Box<Buffered>, _info: &mut GlobalInfo) -> Result<SeekFrom> {
            let len = buf.data().len();
            if 9 > len {
                let needed = 9 - len;
                Err(ErrorKind::MoreDataNeeded(needed).into())
            } else {
                Ok(SeekFrom::Current(9))
            }
        }
        fn read_event(&mut self, buf: &Box<Buffered>) -> Result<(SeekFrom, Event)> {
            let size = 2;
            let len = buf.data().len();
            if size > len {
                Err(ErrorKind::MoreDataNeeded(size - len).into())
            } else {
                println!("{:?}", buf.data());
                match &buf.data()[..2] {
                    b"p1" => {
                        Ok((SeekFrom::Current(3), Event::NewPacket(Packet::new())))
                    },
                    b"e1" => {
                        Ok((SeekFrom::Current(3), Event::MoreDataNeeded(0)))
                    }
                    _ => Err(ErrorKind::InvalidData.into())
                }
            }
        }
    }

    impl Descriptor for DummyDes {
        fn create(&self) -> Box<Demuxer> {
            Box::new(DummyDemuxer {})
        }
        fn describe<'a>(&'a self) -> &'a Descr {
            &self.d
        }
        fn probe(&self, data: &[u8]) -> u8 {
            match data {
                b"dummy" => 100,
                _ => 0,
            }
        }
    }

    const DUMMY_DES: &Descriptor = &DummyDes {
        d: Descr {
            name: "dummy",
            demuxer: "dummy",
            description: "Dummy dem",
            extensions: &["dm", "dum"],
            mime: &["application/dummy"],
        },
    };

    #[test]
    fn probe() {
        let demuxers: &[&'static Descriptor] = &[DUMMY_DES];

        demuxers.probe(b"dummy").unwrap();
    }

    use std::io::Cursor;
    use buffer::*;

    #[test]
    fn read_headers() {
        let buf = b"dummy header";
        let r = AccReader::with_capacity(4, Cursor::new(buf));
        let d = DUMMY_DES.create();
        let mut c = Context::new(d, Box::new(r));

        c.read_headers().unwrap();
    }

    #[test]
    fn read_event() {
        let buf = b"dummy header p1 e1 p1 ";

        let r = AccReader::with_capacity(4, Cursor::new(buf));
        let d = DUMMY_DES.create();
        let mut c = Context::new(d, Box::new(r));

        c.read_headers().unwrap();

        println!("{:?}", c.read_event());
        println!("{:?}", c.read_event());
        println!("{:?}", c.read_event());
        println!("{:?}", c.read_event());
    }
}
<commit_msg>format: Use the stream timebase as default for the packet<commit_after>use data::packet::Packet;
use data::rational::Rational64;
use stream::Stream;
use error::*;

use buffer::Buffered;
use std::io::SeekFrom;

#[derive(Clone, Debug, PartialEq)]
pub struct GlobalInfo {
    pub duration: Option<u64>,
    pub timebase: Option<Rational64>,
    pub streams: Vec<Stream>,
}

#[derive(Clone, Debug, PartialEq)]
pub enum Event {
    NewPacket(Packet),
    NewStream(Stream),
    MoreDataNeeded(usize),
}

pub trait Demuxer : Send {
    fn read_headers(&mut self, buf: &Box<Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom>;
    fn read_event(&mut self, buf: &Box<Buffered>) -> Result<(SeekFrom, Event)>;
}

#[derive(Clone, Debug, PartialEq)]
pub struct Descr {
    pub name: &'static str,
    pub demuxer: &'static str,
    pub description: &'static str,
    pub extensions: &'static [&'static str],
    pub mime: &'static [&'static str],
}

pub trait Descriptor {
    fn create(&self) -> Box<Demuxer>;
    fn describe<'a>(&'a self) -> &'a Descr;

    fn probe(&self, data: &[u8]) -> u8;
}

pub struct Context {
    demuxer: Box<Demuxer>,
    reader: Box<Buffered>,
    pub info: GlobalInfo,
}

impl Context {
    pub fn new<R: Buffered + 'static>(demuxer: Box<Demuxer>, reader: Box<R>) -> Self {
        Context {
            demuxer: demuxer,
            reader: reader,
            info: GlobalInfo {
                duration: None,
                timebase: None,
                streams: Vec::with_capacity(2),
            },
        }
    }

    fn read_headers_internal(&mut self) -> Result<()> {
        let ref mut demux = self.demuxer;
        let res = demux.read_headers(&self.reader, &mut self.info);
        match res {
            Err(e) => Err(e),
            Ok(seek) => {
                //TODO: handle seeking here
                let res = self.reader.seek(seek);
                println!("stream now at index: {:?}", res);
                Ok(())
            }
        }
    }

    pub fn read_headers(&mut self) -> Result<()> {
        loop {
            println!("Filling");
            try!(self.reader.fill_buf());
            match self.read_headers_internal() {
                Err(e) => match e {
                    Error(ErrorKind::MoreDataNeeded(needed), _) => {
                        self.reader.grow(needed);
                    },
                    _ => return Err(e)
                },
                Ok(_) => return Ok(()),
            }
        }
    }

    fn read_event_internal(&mut self) -> Result<Event> {
        let ref mut demux = self.demuxer;

        let res = demux.read_event(&self.reader);
        match res {
            Err(e) => Err(e),
            Ok((seek, mut event)) => {
                //TODO: handle seeking here
                let res = self.reader.seek(seek);
                if let Event::NewStream(ref st) = event {
                    self.info.streams.push(st.clone());
                }
                if let Event::MoreDataNeeded(size) = event {
                    return Err(ErrorKind::MoreDataNeeded(size).into());
                }
                if let Event::NewPacket(ref mut pkt) = event {
                    if pkt.t.timebase.is_none() {
                        if let Some(ref st) = self.info.streams.iter().find(|s| {
                            s.index as isize == pkt.stream_index
                        }) {
                            pkt.t.timebase = Some(st.timebase);
                        }
                    }
                }
                Ok(event)
            }
        }
    }

    pub fn read_event(&mut self) -> Result<Event> {
        // TODO: guard against infiniloops and maybe factor the loop.
        loop {
            match self.read_event_internal() {
                Err(e) => match e {
                    Error(ErrorKind::MoreDataNeeded(needed), _) => {
                        self.reader.grow(needed);
                        try!(self.reader.fill_buf());
                    },
                    _ => return Err(e)
                },
                Ok(ev) => return Ok(ev),
            }
        }
    }
}

pub const PROBE_DATA: usize = 4 * 1024;
pub const PROBE_SCORE_EXTENSION: u8 = 50;

// TODO:
// IntoIterator<Item = &'static Descriptor> is confusing

pub trait Probe {
    fn probe(&self, data: &[u8]) -> Option<&'static Descriptor>;
}

impl<'a> Probe for [&'static Descriptor] {
    fn probe(&self, data: &[u8]) -> Option<&'static Descriptor> {
        let mut max = u8::min_value();
        let mut candidate: Option<&'static Descriptor> = None;
        for desc in self {
            let score = desc.probe(data);

            if score > max {
                max = score;
                candidate = Some(*desc);
            }
        }

        if max > PROBE_SCORE_EXTENSION {
            candidate
        } else {
            None
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::io::SeekFrom;
    use data::packet::Packet;

    struct DummyDes {
        d: Descr,
    }

    struct DummyDemuxer {}

    impl Demuxer for DummyDemuxer {
        fn read_headers(&mut self, buf: &Box<Buffered>, _info: &mut GlobalInfo) -> Result<SeekFrom> {
            let len = buf.data().len();
            if 9 > len {
                let needed = 9 - len;
                Err(ErrorKind::MoreDataNeeded(needed).into())
            } else {
                Ok(SeekFrom::Current(9))
            }
        }
        fn read_event(&mut self, buf: &Box<Buffered>) -> Result<(SeekFrom, Event)> {
            let size = 2;
            let len = buf.data().len();
            if size > len {
                Err(ErrorKind::MoreDataNeeded(size - len).into())
            } else {
                println!("{:?}", buf.data());
                match &buf.data()[..2] {
                    b"p1" => {
                        Ok((SeekFrom::Current(3), Event::NewPacket(Packet::new())))
                    },
                    b"e1" => {
                        Ok((SeekFrom::Current(3), Event::MoreDataNeeded(0)))
                    }
                    _ => Err(ErrorKind::InvalidData.into())
                }
            }
        }
    }

    impl Descriptor for DummyDes {
        fn create(&self) -> Box<Demuxer> {
            Box::new(DummyDemuxer {})
        }
        fn describe<'a>(&'a self) -> &'a Descr {
            &self.d
        }
        fn probe(&self, data: &[u8]) -> u8 {
            match data {
                b"dummy" => 100,
                _ => 0,
            }
        }
    }

    const DUMMY_DES: &Descriptor = &DummyDes {
        d: Descr {
            name: "dummy",
            demuxer: "dummy",
            description: "Dummy dem",
            extensions: &["dm", "dum"],
            mime: &["application/dummy"],
        },
    };

    #[test]
    fn probe() {
        let demuxers: &[&'static Descriptor] = &[DUMMY_DES];

        demuxers.probe(b"dummy").unwrap();
    }

    use std::io::Cursor;
    use buffer::*;

    #[test]
    fn read_headers() {
        let buf = b"dummy header";
        let r = AccReader::with_capacity(4, Cursor::new(buf));
        let d = DUMMY_DES.create();
        let mut c = Context::new(d, Box::new(r));

        c.read_headers().unwrap();
    }

    #[test]
    fn read_event() {
        let buf = b"dummy header p1 e1 p1 ";

        let r = AccReader::with_capacity(4, Cursor::new(buf));
        let d = DUMMY_DES.create();
        let mut c = Context::new(d, Box::new(r));

        c.read_headers().unwrap();

        println!("{:?}", c.read_event());
        println!("{:?}", c.read_event());
        println!("{:?}", c.read_event());
        println!("{:?}", c.read_event());
    }
}
<|endoftext|>