<commit_before>use native;
use std::comm::channel;

struct WorkUnit {
    arg: int,
    callback: Sender<int>,
}

enum MessageToWorker {
    Work(WorkUnit),
    Halt,
}

enum MessageToDispatcher {
    Dispatch(WorkUnit),
    HaltAll,
    RegisterWorker(Sender<Sender<Sender<MessageToWorker>>>),
}

struct WorkQueue {
    dispatcher: Sender<MessageToDispatcher>,
}

impl WorkQueue {
    pub fn create(worker_count: int) -> WorkQueue {
        let (dispatcher, dispatcher_inbox) = channel::<MessageToDispatcher>();

        // dispatcher
        native::task::spawn(proc() {
            let (want_work, idle_worker) = channel::<Sender<MessageToWorker>>();
            let mut worker_count = 0;
            let inbox = dispatcher_inbox;
            let idle_worker = idle_worker;
            loop {
                match inbox.recv() {
                    Dispatch(work_item) => {
                        idle_worker.recv().send(Work(work_item));
                    },
                    RegisterWorker(want_idle_sender) => {
                        worker_count += 1;
                        want_idle_sender.send(want_work.clone());
                    }
                    HaltAll => {
                        while worker_count > 0 {
                            idle_worker.recv().send(Halt);
                            worker_count -= 1;
                        }
                        return;
                    },
                };
            }
        });

        for _ in range(0, worker_count) {
            // worker
            let (reg_s, reg_r) = channel::<Sender<Sender<MessageToWorker>>>();
            dispatcher.send(RegisterWorker(reg_s));
            let want_work = reg_r.recv();

            native::task::spawn(proc() {
                let want_work = want_work;
                loop {
                    let (reply_with_work, get_work_unit) = channel::<MessageToWorker>();
                    want_work.send(reply_with_work);
                    let work_unit = match get_work_unit.recv() {
                        Work(wu) => wu,
                        Halt     => return
                    };
                    let rv = work_unit.arg * 2;
                    work_unit.callback.send(rv);
                }
            });
        }

        return WorkQueue{dispatcher: dispatcher};
    }

    pub fn execute(&self, arg: int) -> Receiver<int> {
        let (reply_with_rv, wait_for_rv) = channel::<int>();
        self.dispatcher.send(Dispatch(WorkUnit{arg: arg, callback: reply_with_rv}));
        return wait_for_rv;
    }
}

impl Drop for WorkQueue {
    fn drop(&mut self) {
        self.dispatcher.send(HaltAll);
    }
}

#[test]
fn test_queue() {
    let queue = WorkQueue::create(3);
    let mut promise_list: ~[Receiver<int>] = ~[];
    for c in range(0, 10) {
        let rv = queue.execute(c);
        promise_list.push(rv);
    }
    let return_list = promise_list.map(|promise| promise.recv());
    assert_eq!(return_list, ~[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]);
}
<commit_msg>create workers on demand<commit_after>use native;
use std::comm::channel;

struct WorkUnit {
    arg: int,
    callback: Sender<int>,
}

enum MessageToWorker {
    Work(WorkUnit),
    Halt,
}

enum MessageToDispatcher {
    Dispatch(WorkUnit),
    HaltAll,
    RegisterWorker(Sender<Sender<Sender<MessageToWorker>>>),
}

struct WorkQueue {
    dispatcher: Sender<MessageToDispatcher>,
}

impl WorkQueue {
    pub fn create() -> WorkQueue {
        let (dispatcher, dispatcher_inbox) = channel::<MessageToDispatcher>();

        // dispatcher
        native::task::spawn(proc() {
            let (want_work, idle_worker) = channel::<Sender<MessageToWorker>>();
            let mut worker_count = 0;
            let inbox = dispatcher_inbox;
            let idle_worker = idle_worker;
            loop {
                match inbox.recv() {
                    Dispatch(work_item) => {
                        idle_worker.recv().send(Work(work_item));
                    },
                    RegisterWorker(want_idle_sender) => {
                        worker_count += 1;
                        want_idle_sender.send(want_work.clone());
                    }
                    HaltAll => {
                        while worker_count > 0 {
                            idle_worker.recv().send(Halt);
                            worker_count -= 1;
                        }
                        return;
                    },
                };
            }
        });
        return WorkQueue{dispatcher: dispatcher};
    }

    pub fn spawn_worker(&self) {
        let (reg_s, reg_r) = channel::<Sender<Sender<MessageToWorker>>>();
        self.dispatcher.send(RegisterWorker(reg_s));
        let want_work = reg_r.recv();

        native::task::spawn(proc() {
            let want_work = want_work;
            loop {
                let (reply_with_work, get_work_unit) = channel::<MessageToWorker>();
                want_work.send(reply_with_work);
                let work_unit = match get_work_unit.recv() {
                    Work(wu) => wu,
                    Halt     => return
                };
                let rv = work_unit.arg * 2;
                work_unit.callback.send(rv);
            }
        });
    }

    pub fn execute(&self, arg: int) -> Receiver<int> {
        let (reply_with_rv, wait_for_rv) = channel::<int>();
        self.dispatcher.send(Dispatch(WorkUnit{arg: arg, callback: reply_with_rv}));
        return wait_for_rv;
    }
}

impl Drop for WorkQueue {
    fn drop(&mut self) {
        self.dispatcher.send(HaltAll);
    }
}

#[test]
fn test_queue() {
    let queue = WorkQueue::create();
    for _ in range(0, 3) {
        queue.spawn_worker();
    }
    let mut promise_list: ~[Receiver<int>] = ~[];
    for c in range(0, 10) {
        let rv = queue.execute(c);
        promise_list.push(rv);
    }
    let return_list = promise_list.map(|promise| promise.recv());
    assert_eq!(return_list, ~[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]);
}
<|endoftext|>


<commit_before>// pest. The Elegant Parser
// Copyright (C) 2017  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::rc::Rc;

use super::error::Error;
use super::inputs::{Input, Position, Span};
use super::inputs::position;
use super::iterators::{pairs, QueueableToken};
use super::RuleType;

/// An `enum` specifying the current lookahead status of a `ParserState`.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Lookahead {
    Positive,
    Negative,
    None
}

/// A `struct` which contains the complete state of a `Parser`.
#[derive(Debug)]
pub struct ParserState<R: RuleType, I: Input> {
    queue: Vec<QueueableToken<R>>,
    lookahead: Lookahead,
    is_atomic: bool,
    pos_attempts: Vec<R>,
    neg_attempts: Vec<R>,
    attempt_pos: usize,
    /// Stack of `Span`s
    pub stack: Vec<Span<I>>
}

/// Creates a `ParserState` from an `Input`, supplying it to a closure `f`.
///
/// # Examples
///
/// ```
/// # use std::rc::Rc;
/// # use pest;
/// # use pest::inputs::StringInput;
///
/// let input = Rc::new(StringInput::new("".to_owned()));
/// pest::state::<(), _, _>(input, |_, pos| {
///     Ok(pos)
/// }).unwrap();
/// ```
pub fn state<R: RuleType, I: Input, F>(
    input: Rc<I>,
    f: F
) -> Result<pairs::Pairs<R, I>, Error<R, I>>
    where
        F: FnOnce(&mut ParserState<R, I>, Position<I>) -> Result<Position<I>, Position<I>>
{
    let mut state = ParserState {
        queue: vec![],
        lookahead: Lookahead::None,
        is_atomic: false,
        pos_attempts: vec![],
        neg_attempts: vec![],
        attempt_pos: 0,
        stack: vec![]
    };

    if f(&mut state, Position::from_start(input.clone())).is_ok() {
        let len = state.queue.len();
        Ok(pairs::new(Rc::new(state.queue), input, 0, len))
    } else {
        Err(Error::ParsingError {
            positives: state.pos_attempts,
            negatives: state.neg_attempts,
            pos: position::new(input, state.attempt_pos)
        })
    }
}

impl<R: RuleType, I: Input> ParserState<R, I> {
    /// Wrapper needed to generate tokens.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::rc::Rc;
    /// # use pest;
    /// # use pest::inputs::StringInput;
    /// # #[allow(non_camel_case_types)]
    /// # #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
    /// enum Rule {
    ///     a
    /// }
    ///
    /// let input = Rc::new(StringInput::new("a".to_owned()));
    /// let pairs: Vec<_> = pest::state(input, |state, pos| {
    ///     state.rule(Rule::a, pos, |_, p| Ok(p))
    /// }).unwrap().collect();
    ///
    /// assert_eq!(pairs.len(), 1);
    /// ```
    #[inline]
    pub fn rule<F>(&mut self, rule: R, pos: Position<I>, f: F) -> Result<Position<I>, Position<I>>
    where
        F: FnOnce(&mut ParserState<R, I>, Position<I>) -> Result<Position<I>, Position<I>>
    {
        let actual_pos = pos.pos();
        let index = self.queue.len();

        self.track(rule, actual_pos);

        if self.lookahead == Lookahead::None && !self.is_atomic {
            // Pair's position will only be known after running the closure.
            self.queue.push(QueueableToken::Start { pair: 0, pos: actual_pos });
        }

        let result = f(self, pos);

        if self.lookahead == Lookahead::None && !self.is_atomic {
            if let Ok(ref pos) = result {
                // Storing the pair's index in the first token that was added before the closure was
                // run.
                let new_index = self.queue.len();
                match self.queue[index] {
                    QueueableToken::Start { ref mut pair, .. } => *pair = new_index,
                    _ => unreachable!()
                };

                self.queue.push(QueueableToken::End { rule: rule, pos: pos.pos() });
            } else {
                self.queue.truncate(index);
            }
        }

        result
    }

    fn track(&mut self, rule: R, pos: usize) {
        if self.is_atomic {
            return;
        }

        let mut attempts = if self.lookahead != Lookahead::Negative {
            &mut self.pos_attempts
        } else {
            &mut self.neg_attempts
        };

        if pos > self.attempt_pos {
            attempts.clear();
            self.attempt_pos = pos;
        }

        attempts.push(rule);
    }

    /// Wrapper which removes `Tokens` in case of a sequence's failure.
    ///
    /// Usually used in conjunction with
    /// [`Position::sequence`](../inputs/struct.Position#method.sequence).
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::rc::Rc;
    /// # use pest;
    /// # use pest::inputs::StringInput;
    /// # #[allow(non_camel_case_types)]
    /// # #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
    /// enum Rule {
    ///     a
    /// }
    ///
    /// let input = Rc::new(StringInput::new("a".to_owned()));
    /// let pairs: Vec<_> = pest::state(input, |state, pos| {
    ///     state.sequence(move |state| {
    ///         pos.sequence(|p| {
    ///             state.rule(Rule::a, p, |_, p| Ok(p)).and_then(|p| {
    ///                 p.match_string("b")
    ///             })
    ///         })
    ///     }).or_else(|p| {
    ///         Ok(p)
    ///     })
    /// }).unwrap().collect();
    ///
    /// assert_eq!(pairs.len(), 0);
    /// ```
    #[inline]
    pub fn sequence<F>(&mut self, f: F) -> Result<Position<I>, Position<I>>
    where
        F: FnOnce(&mut ParserState<R, I>) -> Result<Position<I>, Position<I>>
    {
        let index = self.queue.len();

        let result = f(self);

        if result.is_err() {
            self.queue.truncate(index);
        }

        result
    }

    /// Wrapper which stops `Token`s from being generated.
    ///
    /// Usually used in conjunction with
    /// [`Position::lookahead`](../inputs/struct.Position#method.lookahead).
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::rc::Rc;
    /// # use pest;
    /// # use pest::inputs::StringInput;
    /// # #[allow(non_camel_case_types)]
    /// # #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
    /// enum Rule {
    ///     a
    /// }
    ///
    /// let input = Rc::new(StringInput::new("a".to_owned()));
    /// let pairs: Vec<_> = pest::state(input, |state, pos| {
    ///     state.lookahead(true, move |state| {
    ///         state.rule(Rule::a, pos, |_, p| Ok(p))
    ///     })
    /// }).unwrap().collect();
    ///
    /// assert_eq!(pairs.len(), 0);
    /// ```
    #[inline]
    pub fn lookahead<F>(&mut self, is_positive: bool, f: F) -> Result<Position<I>, Position<I>>
    where
        F: FnOnce(&mut ParserState<R, I>) -> Result<Position<I>, Position<I>>
    {
        let initial_lookahead = self.lookahead;

        self.lookahead = if is_positive {
            Lookahead::Positive
        } else {
            Lookahead::Negative
        };

        let result = f(self);

        self.lookahead = initial_lookahead;

        result
    }

    /// Wrapper which stops `Token`s from being generated according to `is_atomic`.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::rc::Rc;
    /// # use pest;
    /// # use pest::inputs::StringInput;
    /// # #[allow(non_camel_case_types)]
    /// # #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
    /// enum Rule {
    ///     a
    /// }
    ///
    /// let input = Rc::new(StringInput::new("a".to_owned()));
    /// let pairs: Vec<_> = pest::state(input, |state, pos| {
    ///     state.atomic(true, move |state| {
    ///         state.rule(Rule::a, pos, |_, p| Ok(p))
    ///     })
    /// }).unwrap().collect();
    ///
    /// assert_eq!(pairs.len(), 0);
    /// ```
    #[inline]
    pub fn atomic<F>(&mut self, is_atomic: bool, f: F) -> Result<Position<I>, Position<I>>
    where
        F: FnOnce(&mut ParserState<R, I>) -> Result<Position<I>, Position<I>>
    {
        let should_toggle = self.is_atomic != is_atomic;

        if should_toggle {
            self.is_atomic = is_atomic;
        }

        let result = f(self);

        if should_toggle {
            self.is_atomic = !is_atomic;
        }

        result
    }

    /// Returns whether the `ParserState` is in an atomic state.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::rc::Rc;
    /// # use pest;
    /// # use pest::inputs::StringInput;
    /// # #[allow(non_camel_case_types)]
    /// # #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
    /// enum Rule {
    ///     a
    /// }
    ///
    /// let input = Rc::new(StringInput::new("a".to_owned()));
    /// let pairs: Vec<_> = pest::state::<(), _, _>(input, |state, pos| {
    ///     assert!(!state.is_atomic());
    ///
    ///     state.atomic(true, move |state| {
    ///         assert!(state.is_atomic());
    ///
    ///         Ok(pos)
    ///     })
    /// }).unwrap().collect();
    ///
    /// assert_eq!(pairs.len(), 0);
    /// ```
    #[inline]
    pub fn is_atomic(&self) -> bool {
        self.is_atomic
    }
}
<commit_msg>Changed attempt tracking.<commit_after>// pest. The Elegant Parser
// Copyright (C) 2017  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::rc::Rc;

use super::error::Error;
use super::inputs::{Input, Position, Span};
use super::inputs::position;
use super::iterators::{pairs, QueueableToken};
use super::RuleType;

/// An `enum` specifying the current lookahead status of a `ParserState`.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Lookahead {
    Positive,
    Negative,
    None
}

/// A `struct` which contains the complete state of a `Parser`.
#[derive(Debug)]
pub struct ParserState<R: RuleType, I: Input> {
    queue: Vec<QueueableToken<R>>,
    lookahead: Lookahead,
    is_atomic: bool,
    pos_attempts: Vec<R>,
    neg_attempts: Vec<R>,
    attempt_pos: usize,
    /// Stack of `Span`s
    pub stack: Vec<Span<I>>
}

/// Creates a `ParserState` from an `Input`, supplying it to a closure `f`.
///
/// # Examples
///
/// ```
/// # use std::rc::Rc;
/// # use pest;
/// # use pest::inputs::StringInput;
///
/// let input = Rc::new(StringInput::new("".to_owned()));
/// pest::state::<(), _, _>(input, |_, pos| {
///     Ok(pos)
/// }).unwrap();
/// ```
pub fn state<R: RuleType, I: Input, F>(
    input: Rc<I>,
    f: F
) -> Result<pairs::Pairs<R, I>, Error<R, I>>
where
    F: FnOnce(&mut ParserState<R, I>, Position<I>) -> Result<Position<I>, Position<I>>
{
    let mut state = ParserState {
        queue: vec![],
        lookahead: Lookahead::None,
        is_atomic: false,
        pos_attempts: vec![],
        neg_attempts: vec![],
        attempt_pos: 0,
        stack: vec![]
    };

    if f(&mut state, Position::from_start(input.clone())).is_ok() {
        let len = state.queue.len();
        Ok(pairs::new(Rc::new(state.queue), input, 0, len))
    } else {
        state.pos_attempts.sort();
        state.pos_attempts.dedup();
        state.neg_attempts.sort();
        state.neg_attempts.dedup();

        Err(Error::ParsingError {
            positives: state.pos_attempts,
            negatives: state.neg_attempts,
            pos: position::new(input, state.attempt_pos)
        })
    }
}

impl<R: RuleType, I: Input> ParserState<R, I> {
    /// Wrapper needed to generate tokens.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::rc::Rc;
    /// # use pest;
    /// # use pest::inputs::StringInput;
    /// # #[allow(non_camel_case_types)]
    /// # #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
    /// enum Rule {
    ///     a
    /// }
    ///
    /// let input = Rc::new(StringInput::new("a".to_owned()));
    /// let pairs: Vec<_> = pest::state(input, |state, pos| {
    ///     state.rule(Rule::a, pos, |_, p| Ok(p))
    /// }).unwrap().collect();
    ///
    /// assert_eq!(pairs.len(), 1);
    /// ```
    #[inline]
    pub fn rule<F>(&mut self, rule: R, pos: Position<I>, f: F) -> Result<Position<I>, Position<I>>
    where
        F: FnOnce(&mut ParserState<R, I>, Position<I>) -> Result<Position<I>, Position<I>>
    {
        let actual_pos = pos.pos();
        let index = self.queue.len();

        let (pos_attempts_index, neg_attempts_index) = if actual_pos == self.attempt_pos {
            (self.pos_attempts.len(), self.neg_attempts.len())
        } else { // Attempts have not been cleared yet since the attempt_pos is older.
            (0, 0)
        };

        if self.lookahead == Lookahead::None && !self.is_atomic {
            // Pair's position will only be known after running the closure.
            self.queue.push(QueueableToken::Start { pair: 0, pos: actual_pos });
        }

        let result = f(self, pos);

        if result.is_err() {
            self.track(rule, actual_pos, pos_attempts_index, neg_attempts_index);
        }

        if self.lookahead == Lookahead::None && !self.is_atomic {
            if let Ok(ref pos) = result {
                // Storing the pair's index in the first token that was added before the closure was
                // run.
                let new_index = self.queue.len();
                match self.queue[index] {
                    QueueableToken::Start { ref mut pair, .. } => *pair = new_index,
                    _ => unreachable!()
                };

                self.queue.push(QueueableToken::End { rule: rule, pos: pos.pos() });
            } else {
                self.queue.truncate(index);
            }
        }

        result
    }

    fn track(&mut self, rule: R, pos: usize, pos_attempts_index: usize, neg_attempts_index: usize) {
        if self.is_atomic {
            return;
        }

        // If nested rules made no progress, there is no use to report them; it's only useful to
        // track the current rule.
        if pos == self.attempt_pos {
            self.pos_attempts.truncate(pos_attempts_index);
            self.neg_attempts.truncate(neg_attempts_index);
        }println!("{} {} {:?} {:?} {}", pos, self.attempt_pos, rule, self.pos_attempts, pos_attempts_index);

        let mut attempts = if self.lookahead != Lookahead::Negative {
            &mut self.pos_attempts
        } else {
            &mut self.neg_attempts
        };

        if pos > self.attempt_pos {
            attempts.clear();
            self.attempt_pos = pos;
        }

        if pos == self.attempt_pos {
            attempts.push(rule);
        }
    }

    /// Wrapper which removes `Tokens` in case of a sequence's failure.
    ///
    /// Usually used in conjunction with
    /// [`Position::sequence`](../inputs/struct.Position#method.sequence).
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::rc::Rc;
    /// # use pest;
    /// # use pest::inputs::StringInput;
    /// # #[allow(non_camel_case_types)]
    /// # #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
    /// enum Rule {
    ///     a
    /// }
    ///
    /// let input = Rc::new(StringInput::new("a".to_owned()));
    /// let pairs: Vec<_> = pest::state(input, |state, pos| {
    ///     state.sequence(move |state| {
    ///         pos.sequence(|p| {
    ///             state.rule(Rule::a, p, |_, p| Ok(p)).and_then(|p| {
    ///                 p.match_string("b")
    ///             })
    ///         })
    ///     }).or_else(|p| {
    ///         Ok(p)
    ///     })
    /// }).unwrap().collect();
    ///
    /// assert_eq!(pairs.len(), 0);
    /// ```
    #[inline]
    pub fn sequence<F>(&mut self, f: F) -> Result<Position<I>, Position<I>>
    where
        F: FnOnce(&mut ParserState<R, I>) -> Result<Position<I>, Position<I>>
    {
        let index = self.queue.len();

        let result = f(self);

        if result.is_err() {
            self.queue.truncate(index);
        }

        result
    }

    /// Wrapper which stops `Token`s from being generated.
    ///
    /// Usually used in conjunction with
    /// [`Position::lookahead`](../inputs/struct.Position#method.lookahead).
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::rc::Rc;
    /// # use pest;
    /// # use pest::inputs::StringInput;
    /// # #[allow(non_camel_case_types)]
    /// # #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
    /// enum Rule {
    ///     a
    /// }
    ///
    /// let input = Rc::new(StringInput::new("a".to_owned()));
    /// let pairs: Vec<_> = pest::state(input, |state, pos| {
    ///     state.lookahead(true, move |state| {
    ///         state.rule(Rule::a, pos, |_, p| Ok(p))
    ///     })
    /// }).unwrap().collect();
    ///
    /// assert_eq!(pairs.len(), 0);
    /// ```
    #[inline]
    pub fn lookahead<F>(&mut self, is_positive: bool, f: F) -> Result<Position<I>, Position<I>>
    where
        F: FnOnce(&mut ParserState<R, I>) -> Result<Position<I>, Position<I>>
    {
        let initial_lookahead = self.lookahead;

        self.lookahead = if is_positive {
            Lookahead::Positive
        } else {
            Lookahead::Negative
        };

        let result = f(self);

        self.lookahead = initial_lookahead;

        result
    }

    /// Wrapper which stops `Token`s from being generated according to `is_atomic`.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::rc::Rc;
    /// # use pest;
    /// # use pest::inputs::StringInput;
    /// # #[allow(non_camel_case_types)]
    /// # #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
    /// enum Rule {
    ///     a
    /// }
    ///
    /// let input = Rc::new(StringInput::new("a".to_owned()));
    /// let pairs: Vec<_> = pest::state(input, |state, pos| {
    ///     state.atomic(true, move |state| {
    ///         state.rule(Rule::a, pos, |_, p| Ok(p))
    ///     })
    /// }).unwrap().collect();
    ///
    /// assert_eq!(pairs.len(), 0);
    /// ```
    #[inline]
    pub fn atomic<F>(&mut self, is_atomic: bool, f: F) -> Result<Position<I>, Position<I>>
    where
        F: FnOnce(&mut ParserState<R, I>) -> Result<Position<I>, Position<I>>
    {
        let should_toggle = self.is_atomic != is_atomic;

        if should_toggle {
            self.is_atomic = is_atomic;
        }

        let result = f(self);

        if should_toggle {
            self.is_atomic = !is_atomic;
        }

        result
    }

    /// Returns whether the `ParserState` is in an atomic state.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::rc::Rc;
    /// # use pest;
    /// # use pest::inputs::StringInput;
    /// # #[allow(non_camel_case_types)]
    /// # #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
    /// enum Rule {
    ///     a
    /// }
    ///
    /// let input = Rc::new(StringInput::new("a".to_owned()));
    /// let pairs: Vec<_> = pest::state::<(), _, _>(input, |state, pos| {
    ///     assert!(!state.is_atomic());
    ///
    ///     state.atomic(true, move |state| {
    ///         assert!(state.is_atomic());
    ///
    ///         Ok(pos)
    ///     })
    /// }).unwrap().collect();
    ///
    /// assert_eq!(pairs.len(), 0);
    /// ```
    #[inline]
    pub fn is_atomic(&self) -> bool {
        self.is_atomic
    }
}
<|endoftext|>


<commit_before>//! Renderer for function calls.

use hir::{AsAssocItem, HasSource, HirDisplay};
use ide_db::SymbolKind;
use itertools::Itertools;
use syntax::ast;

use crate::{
    item::{CompletionItem, CompletionItemKind, CompletionKind, CompletionRelevance, ImportEdit},
    render::{
        builder_ext::Params, compute_exact_name_match, compute_ref_match, compute_type_match,
        RenderContext,
    },
};

pub(crate) fn render_fn(
    ctx: RenderContext<'_>,
    import_to_add: Option<ImportEdit>,
    local_name: Option<hir::Name>,
    fn_: hir::Function,
) -> Option<CompletionItem> {
    let _p = profile::span("render_fn");
    Some(FunctionRender::new(ctx, None, local_name, fn_, false)?.render(import_to_add))
}

pub(crate) fn render_method(
    ctx: RenderContext<'_>,
    import_to_add: Option<ImportEdit>,
    receiver: Option<hir::Name>,
    local_name: Option<hir::Name>,
    fn_: hir::Function,
) -> Option<CompletionItem> {
    let _p = profile::span("render_method");
    Some(FunctionRender::new(ctx, receiver, local_name, fn_, true)?.render(import_to_add))
}

#[derive(Debug)]
struct FunctionRender<'a> {
    ctx: RenderContext<'a>,
    name: String,
    receiver: Option<hir::Name>,
    func: hir::Function,
    ast_node: ast::Fn,
    is_method: bool,
}

impl<'a> FunctionRender<'a> {
    fn new(
        ctx: RenderContext<'a>,
        receiver: Option<hir::Name>,
        local_name: Option<hir::Name>,
        fn_: hir::Function,
        is_method: bool,
    ) -> Option<FunctionRender<'a>> {
        let name = local_name.unwrap_or_else(|| fn_.name(ctx.db())).to_string();
        let ast_node = fn_.source(ctx.db())?.value;

        Some(FunctionRender { ctx, name, receiver, func: fn_, ast_node, is_method })
    }

    fn render(self, import_to_add: Option<ImportEdit>) -> CompletionItem {
        let params = self.params();
        let call = if let Some(receiver) = &self.receiver {
            format!("{}.{}", receiver, &self.name)
        } else {
            self.name.clone()
        };
        let mut item =
            CompletionItem::new(CompletionKind::Reference, self.ctx.source_range(), call.clone());
        item.kind(self.kind())
            .set_documentation(self.ctx.docs(self.func))
            .set_deprecated(
                self.ctx.is_deprecated(self.func) || self.ctx.is_deprecated_assoc_item(self.func),
            )
            .detail(self.detail())
            .add_call_parens(self.ctx.completion, call.clone(), params);

        if import_to_add.is_none() {
            let db = self.ctx.db();
            if let Some(actm) = self.func.as_assoc_item(db) {
                if let Some(trt) = actm.containing_trait_or_trait_impl(db) {
                    item.trait_name(trt.name(db).to_string());
                }
            }
        }

        item.add_import(import_to_add).lookup_by(self.name);

        let ret_type = self.func.ret_type(self.ctx.db());
        item.set_relevance(CompletionRelevance {
            type_match: compute_type_match(self.ctx.completion, &ret_type),
            exact_name_match: compute_exact_name_match(self.ctx.completion, &call),
            ..CompletionRelevance::default()
        });

        if let Some(ref_match) = compute_ref_match(self.ctx.completion, &ret_type) {
            // FIXME
            // For now we don't properly calculate the edits for ref match
            // completions on methods, so we've disabled them. See #8058.
            if !self.is_method {
                item.ref_match(ref_match);
            }
        }

        item.build()
    }

    fn detail(&self) -> String {
        let ret_ty = self.func.ret_type(self.ctx.db());
        let ret = if ret_ty.is_unit() {
            // Omit the return type if it is the unit type
            String::new()
        } else {
            format!(" {}", self.ty_display())
        };

        format!("fn({}){}", self.params_display(), ret)
    }

    fn params_display(&self) -> String {
        if let Some(self_param) = self.func.self_param(self.ctx.db()) {
            let params = self
                .func
                .assoc_fn_params(self.ctx.db())
                .into_iter()
                .skip(1) // skip the self param because we are manually handling that
                .map(|p| p.ty().display(self.ctx.db()).to_string());

            std::iter::once(self_param.display(self.ctx.db()).to_owned()).chain(params).join(", ")
        } else {
            let params = self
                .func
                .assoc_fn_params(self.ctx.db())
                .into_iter()
                .map(|p| p.ty().display(self.ctx.db()).to_string())
                .join(", ");
            params
        }
    }

    fn ty_display(&self) -> String {
        let ret_ty = self.func.ret_type(self.ctx.db());

        format!("-> {}", ret_ty.display(self.ctx.db()))
    }

    fn add_arg(&self, arg: &str, ty: &hir::Type) -> String {
        if let Some(derefed_ty) = ty.remove_ref() {
            for (name, local) in self.ctx.completion.locals.iter() {
                if name == arg && local.ty(self.ctx.db()) == derefed_ty {
                    let mutability = if ty.is_mutable_reference() { "&mut " } else { "&" };
                    return format!("{}{}", mutability, arg);
                }
            }
        }
        arg.to_string()
    }

    fn params(&self) -> Params {
        let ast_params = match self.ast_node.param_list() {
            Some(it) => it,
            None => return Params::Named(Vec::new()),
        };

        let mut params_pats = Vec::new();
        let params_ty = if self.ctx.completion.has_dot_receiver() || self.receiver.is_some() {
            self.func.method_params(self.ctx.db()).unwrap_or_default()
        } else {
            if let Some(s) = ast_params.self_param() {
                cov_mark::hit!(parens_for_method_call_as_assoc_fn);
                params_pats.push(Some(s.to_string()));
            }
            self.func.assoc_fn_params(self.ctx.db())
        };
        params_pats
            .extend(ast_params.params().into_iter().map(|it| it.pat().map(|it| it.to_string())));

        let params = params_pats
            .into_iter()
            .zip(params_ty)
            .flat_map(|(pat, param_ty)| {
                let pat = pat?;
                let name = pat;
                let arg = name.trim_start_matches("mut ").trim_start_matches('_');
                Some(self.add_arg(arg, param_ty.ty()))
            })
            .collect();
        Params::Named(params)
    }

    fn kind(&self) -> CompletionItemKind {
        if self.func.self_param(self.ctx.db()).is_some() {
            CompletionItemKind::Method
        } else {
            SymbolKind::Function.into()
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        tests::{check_edit, check_edit_with_config, TEST_CONFIG},
        CompletionConfig,
    };

    #[test]
    fn inserts_parens_for_function_calls() {
        cov_mark::check!(inserts_parens_for_function_calls);
        check_edit(
            "no_args",
            r#"
fn no_args() {}
fn main() { no_$0 }
"#,
            r#"
fn no_args() {}
fn main() { no_args()$0 }
"#,
        );

        check_edit(
            "with_args",
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_$0 }
"#,
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_args(${1:x}, ${2:y})$0 }
"#,
        );

        check_edit(
            "foo",
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn bar(s: &S) { s.f$0 }
"#,
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn bar(s: &S) { s.foo()$0 }
"#,
        );

        check_edit(
            "foo",
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {}
}
fn bar(s: &S) {
    s.f$0
}
"#,
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {}
}
fn bar(s: &S) {
    s.foo(${1:x})$0
}
"#,
        );

        check_edit(
            "foo",
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {
        $0
    }
}
"#,
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {
        self.foo(${1:x})$0
    }
}
"#,
        );
    }

    #[test]
    fn parens_for_method_call_as_assoc_fn() {
        cov_mark::check!(parens_for_method_call_as_assoc_fn);
        check_edit(
            "foo",
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn main() { S::f$0 }
"#,
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn main() { S::foo(${1:&self})$0 }
"#,
        );
    }

    #[test]
    fn suppress_arg_snippets() {
        cov_mark::check!(suppress_arg_snippets);
        check_edit_with_config(
            CompletionConfig { add_call_argument_snippets: false, ..TEST_CONFIG },
            "with_args",
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_$0 }
"#,
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_args($0) }
"#,
        );
    }

    #[test]
    fn strips_underscores_from_args() {
        check_edit(
            "foo",
            r#"
fn foo(_foo: i32, ___bar: bool, ho_ge_: String) {}
fn main() { f$0 }
"#,
            r#"
fn foo(_foo: i32, ___bar: bool, ho_ge_: String) {}
fn main() { foo(${1:foo}, ${2:bar}, ${3:ho_ge_})$0 }
"#,
        );
    }

    #[test]
    fn insert_ref_when_matching_local_in_scope() {
        check_edit(
            "ref_arg",
            r#"
struct Foo {}
fn ref_arg(x: &Foo) {}
fn main() {
    let x = Foo {};
    ref_ar$0
}
"#,
            r#"
struct Foo {}
fn ref_arg(x: &Foo) {}
fn main() {
    let x = Foo {};
    ref_arg(${1:&x})$0
}
"#,
        );
    }

    #[test]
    fn insert_mut_ref_when_matching_local_in_scope() {
        check_edit(
            "ref_arg",
            r#"
struct Foo {}
fn ref_arg(x: &mut Foo) {}
fn main() {
    let x = Foo {};
    ref_ar$0
}
"#,
            r#"
struct Foo {}
fn ref_arg(x: &mut Foo) {}
fn main() {
    let x = Foo {};
    ref_arg(${1:&mut x})$0
}
"#,
        );
    }

    #[test]
    fn insert_ref_when_matching_local_in_scope_for_method() {
        check_edit(
            "apply_foo",
            r#"
struct Foo {}
struct Bar {}
impl Bar {
    fn apply_foo(&self, x: &Foo) {}
}

fn main() {
    let x = Foo {};
    let y = Bar {};
    y.$0
}
"#,
            r#"
struct Foo {}
struct Bar {}
impl Bar {
    fn apply_foo(&self, x: &Foo) {}
}

fn main() {
    let x = Foo {};
    let y = Bar {};
    y.apply_foo(${1:&x})$0
}
"#,
        );
    }

    #[test]
    fn trim_mut_keyword_in_func_completion() {
        check_edit(
            "take_mutably",
            r#"
fn take_mutably(mut x: &i32) {}

fn main() {
    take_m$0
}
"#,
            r#"
fn take_mutably(mut x: &i32) {}

fn main() {
    take_mutably(${1:x})$0
}
"#,
        );
    }
}
<commit_msg>internal: document query implication of completion rendering<commit_after>//! Renderer for function calls.

use hir::{AsAssocItem, HasSource, HirDisplay};
use ide_db::SymbolKind;
use itertools::Itertools;
use syntax::ast;

use crate::{
    item::{CompletionItem, CompletionItemKind, CompletionKind, CompletionRelevance, ImportEdit},
    render::{
        builder_ext::Params, compute_exact_name_match, compute_ref_match, compute_type_match,
        RenderContext,
    },
};

pub(crate) fn render_fn(
    ctx: RenderContext<'_>,
    import_to_add: Option<ImportEdit>,
    local_name: Option<hir::Name>,
    fn_: hir::Function,
) -> Option<CompletionItem> {
    let _p = profile::span("render_fn");
    Some(FunctionRender::new(ctx, None, local_name, fn_, false)?.render(import_to_add))
}

pub(crate) fn render_method(
    ctx: RenderContext<'_>,
    import_to_add: Option<ImportEdit>,
    receiver: Option<hir::Name>,
    local_name: Option<hir::Name>,
    fn_: hir::Function,
) -> Option<CompletionItem> {
    let _p = profile::span("render_method");
    Some(FunctionRender::new(ctx, receiver, local_name, fn_, true)?.render(import_to_add))
}

#[derive(Debug)]
struct FunctionRender<'a> {
    ctx: RenderContext<'a>,
    name: String,
    receiver: Option<hir::Name>,
    func: hir::Function,
    /// NB: having `ast::Fn` here might or might not be a good idea. The problem
    /// with it is that, to get an `ast::`, you want to parse the corresponding
    /// source file. So, when flyimport completions suggest a bunch of
    /// functions, we spend quite some time parsing many files.
    ///
    /// We need ast because we want to access parameter names (patterns). We can
    /// add them to the hir of the function itself, but parameter names are not
    /// something hir cares otherwise.
    ///
    /// Alternatively we can reconstruct params from the function body, but that
    /// would require parsing anyway.
    ///
    /// It seems that just using `ast` is the best choice -- most of parses
    /// should be cached anyway.
    ast_node: ast::Fn,
    is_method: bool,
}

impl<'a> FunctionRender<'a> {
    fn new(
        ctx: RenderContext<'a>,
        receiver: Option<hir::Name>,
        local_name: Option<hir::Name>,
        fn_: hir::Function,
        is_method: bool,
    ) -> Option<FunctionRender<'a>> {
        let name = local_name.unwrap_or_else(|| fn_.name(ctx.db())).to_string();
        let ast_node = fn_.source(ctx.db())?.value;

        Some(FunctionRender { ctx, name, receiver, func: fn_, ast_node, is_method })
    }

    fn render(self, import_to_add: Option<ImportEdit>) -> CompletionItem {
        let params = self.params();
        let call = if let Some(receiver) = &self.receiver {
            format!("{}.{}", receiver, &self.name)
        } else {
            self.name.clone()
        };
        let mut item =
            CompletionItem::new(CompletionKind::Reference, self.ctx.source_range(), call.clone());
        item.kind(self.kind())
            .set_documentation(self.ctx.docs(self.func))
            .set_deprecated(
                self.ctx.is_deprecated(self.func) || self.ctx.is_deprecated_assoc_item(self.func),
            )
            .detail(self.detail())
            .add_call_parens(self.ctx.completion, call.clone(), params);

        if import_to_add.is_none() {
            let db = self.ctx.db();
            if let Some(actm) = self.func.as_assoc_item(db) {
                if let Some(trt) = actm.containing_trait_or_trait_impl(db) {
                    item.trait_name(trt.name(db).to_string());
                }
            }
        }

        item.add_import(import_to_add).lookup_by(self.name);

        let ret_type = self.func.ret_type(self.ctx.db());
        item.set_relevance(CompletionRelevance {
            type_match: compute_type_match(self.ctx.completion, &ret_type),
            exact_name_match: compute_exact_name_match(self.ctx.completion, &call),
            ..CompletionRelevance::default()
        });

        if let Some(ref_match) = compute_ref_match(self.ctx.completion, &ret_type) {
            // FIXME
            // For now we don't properly calculate the edits for ref match
            // completions on methods, so we've disabled them. See #8058.
            if !self.is_method {
                item.ref_match(ref_match);
            }
        }

        item.build()
    }

    fn detail(&self) -> String {
        let ret_ty = self.func.ret_type(self.ctx.db());
        let ret = if ret_ty.is_unit() {
            // Omit the return type if it is the unit type
            String::new()
        } else {
            format!(" {}", self.ty_display())
        };

        format!("fn({}){}", self.params_display(), ret)
    }

    fn params_display(&self) -> String {
        if let Some(self_param) = self.func.self_param(self.ctx.db()) {
            let params = self
                .func
                .assoc_fn_params(self.ctx.db())
                .into_iter()
                .skip(1) // skip the self param because we are manually handling that
                .map(|p| p.ty().display(self.ctx.db()).to_string());

            std::iter::once(self_param.display(self.ctx.db()).to_owned()).chain(params).join(", ")
        } else {
            let params = self
                .func
                .assoc_fn_params(self.ctx.db())
                .into_iter()
                .map(|p| p.ty().display(self.ctx.db()).to_string())
                .join(", ");
            params
        }
    }

    fn ty_display(&self) -> String {
        let ret_ty = self.func.ret_type(self.ctx.db());

        format!("-> {}", ret_ty.display(self.ctx.db()))
    }

    fn add_arg(&self, arg: &str, ty: &hir::Type) -> String {
        if let Some(derefed_ty) = ty.remove_ref() {
            for (name, local) in self.ctx.completion.locals.iter() {
                if name == arg && local.ty(self.ctx.db()) == derefed_ty {
                    let mutability = if ty.is_mutable_reference() { "&mut " } else { "&" };
                    return format!("{}{}", mutability, arg);
                }
            }
        }
        arg.to_string()
    }

    fn params(&self) -> Params {
        let ast_params = match self.ast_node.param_list() {
            Some(it) => it,
            None => return Params::Named(Vec::new()),
        };

        let mut params_pats = Vec::new();
        let params_ty = if self.ctx.completion.has_dot_receiver() || self.receiver.is_some() {
            self.func.method_params(self.ctx.db()).unwrap_or_default()
        } else {
            if let Some(s) = ast_params.self_param() {
                cov_mark::hit!(parens_for_method_call_as_assoc_fn);
                params_pats.push(Some(s.to_string()));
            }
            self.func.assoc_fn_params(self.ctx.db())
        };
        params_pats
            .extend(ast_params.params().into_iter().map(|it| it.pat().map(|it| it.to_string())));

        let params = params_pats
            .into_iter()
            .zip(params_ty)
            .flat_map(|(pat, param_ty)| {
                let pat = pat?;
                let name = pat;
                let arg = name.trim_start_matches("mut ").trim_start_matches('_');
                Some(self.add_arg(arg, param_ty.ty()))
            })
            .collect();
        Params::Named(params)
    }

    fn kind(&self) -> CompletionItemKind {
        if self.func.self_param(self.ctx.db()).is_some() {
            CompletionItemKind::Method
        } else {
            SymbolKind::Function.into()
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        tests::{check_edit, check_edit_with_config, TEST_CONFIG},
        CompletionConfig,
    };

    #[test]
    fn inserts_parens_for_function_calls() {
        cov_mark::check!(inserts_parens_for_function_calls);
        check_edit(
            "no_args",
            r#"
fn no_args() {}
fn main() { no_$0 }
"#,
            r#"
fn no_args() {}
fn main() { no_args()$0 }
"#,
        );

        check_edit(
            "with_args",
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_$0 }
"#,
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_args(${1:x}, ${2:y})$0 }
"#,
        );

        check_edit(
            "foo",
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn bar(s: &S) { s.f$0 }
"#,
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn bar(s: &S) { s.foo()$0 }
"#,
        );

        check_edit(
            "foo",
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {}
}
fn bar(s: &S) {
    s.f$0
}
"#,
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {}
}
fn bar(s: &S) {
    s.foo(${1:x})$0
}
"#,
        );

        check_edit(
            "foo",
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {
        $0
    }
}
"#,
            r#"
struct S {}
impl S {
    fn foo(&self, x: i32) {
        self.foo(${1:x})$0
    }
}
"#,
        );
    }

    #[test]
    fn parens_for_method_call_as_assoc_fn() {
        cov_mark::check!(parens_for_method_call_as_assoc_fn);
        check_edit(
            "foo",
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn main() { S::f$0 }
"#,
            r#"
struct S;
impl S {
    fn foo(&self) {}
}
fn main() { S::foo(${1:&self})$0 }
"#,
        );
    }

    #[test]
    fn suppress_arg_snippets() {
        cov_mark::check!(suppress_arg_snippets);
        check_edit_with_config(
            CompletionConfig { add_call_argument_snippets: false, ..TEST_CONFIG },
            "with_args",
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_$0 }
"#,
            r#"
fn with_args(x: i32, y: String) {}
fn main() { with_args($0) }
"#,
        );
    }

    #[test]
    fn strips_underscores_from_args() {
        check_edit(
            "foo",
            r#"
fn foo(_foo: i32, ___bar: bool, ho_ge_: String) {}
fn main() { f$0 }
"#,
            r#"
fn foo(_foo: i32, ___bar: bool, ho_ge_: String) {}
fn main() { foo(${1:foo}, ${2:bar}, ${3:ho_ge_})$0 }
"#,
        );
    }

    #[test]
    fn insert_ref_when_matching_local_in_scope() {
        check_edit(
            "ref_arg",
            r#"
struct Foo {}
fn ref_arg(x: &Foo) {}
fn main() {
    let x = Foo {};
    ref_ar$0
}
"#,
            r#"
struct Foo {}
fn ref_arg(x: &Foo) {}
fn main() {
    let x = Foo {};
    ref_arg(${1:&x})$0
}
"#,
        );
    }

    #[test]
    fn insert_mut_ref_when_matching_local_in_scope() {
        check_edit(
            "ref_arg",
            r#"
struct Foo {}
fn ref_arg(x: &mut Foo) {}
fn main() {
    let x = Foo {};
    ref_ar$0
}
"#,
            r#"
struct Foo {}
fn ref_arg(x: &mut Foo) {}
fn main() {
    let x = Foo {};
    ref_arg(${1:&mut x})$0
}
"#,
        );
    }

    #[test]
    fn insert_ref_when_matching_local_in_scope_for_method() {
        check_edit(
            "apply_foo",
            r#"
struct Foo {}
struct Bar {}
impl Bar {
    fn apply_foo(&self, x: &Foo) {}
}

fn main() {
    let x = Foo {};
    let y = Bar {};
    y.$0
}
"#,
            r#"
struct Foo {}
struct Bar {}
impl Bar {
    fn apply_foo(&self, x: &Foo) {}
}

fn main() {
    let x = Foo {};
    let y = Bar {};
    y.apply_foo(${1:&x})$0
}
"#,
        );
    }

    #[test]
    fn trim_mut_keyword_in_func_completion() {
        check_edit(
            "take_mutably",
            r#"
fn take_mutably(mut x: &i32) {}

fn main() {
    take_m$0
}
"#,
            r#"
fn take_mutably(mut x: &i32) {}

fn main() {
    take_mutably(${1:x})$0
}
"#,
        );
    }
}
<|endoftext|>


<commit_before>use crate::{t, VERSION};
use std::fmt::Write as _;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::str::FromStr;
use std::{
    env, fmt, fs,
    io::{self, Write},
};

#[derive(Clone, Copy, Eq, PartialEq)]
pub enum Profile {
    Compiler,
    Codegen,
    Library,
    User,
}

impl Profile {
    fn include_path(&self, src_path: &Path) -> PathBuf {
        PathBuf::from(format!("{}/src/bootstrap/defaults/config.{}.toml", src_path.display(), self))
    }

    pub fn all() -> impl Iterator<Item = Self> {
        use Profile::*;
        // N.B. these are ordered by how they are displayed, not alphabetically
        [Library, Compiler, Codegen, User].iter().copied()
    }

    pub fn purpose(&self) -> String {
        use Profile::*;
        match self {
            Library => "Contribute to the standard library",
            Compiler => "Contribute to the compiler or rustdoc",
            Codegen => "Contribute to the compiler, and also modify LLVM or codegen",
            User => "Install Rust from source",
        }
        .to_string()
    }

    pub fn all_for_help(indent: &str) -> String {
        let mut out = String::new();
        for choice in Profile::all() {
            writeln!(&mut out, "{}{}: {}", indent, choice, choice.purpose()).unwrap();
        }
        out
    }
}

impl FromStr for Profile {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "lib" | "library" => Ok(Profile::Library),
            "compiler" | "rustdoc" => Ok(Profile::Compiler),
            "llvm" | "codegen" => Ok(Profile::Codegen),
            "maintainer" | "user" => Ok(Profile::User),
            _ => Err(format!("unknown profile: '{}'", s)),
        }
    }
}

impl fmt::Display for Profile {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Profile::Compiler => write!(f, "compiler"),
            Profile::Codegen => write!(f, "codegen"),
            Profile::Library => write!(f, "library"),
            Profile::User => write!(f, "user"),
        }
    }
}

pub fn setup(src_path: &Path, profile: Profile) {
    let cfg_file = env::var_os("BOOTSTRAP_CONFIG").map(PathBuf::from);

    if cfg_file.as_ref().map_or(false, |f| f.exists()) {
        let file = cfg_file.unwrap();
        println!(
            "error: you asked `x.py` to setup a new config file, but one already exists at `{}`",
            file.display()
        );
        println!("help: try adding `profile = \"{}\"` at the top of {}", profile, file.display());
        println!(
            "note: this will use the configuration in {}",
            profile.include_path(src_path).display()
        );
        std::process::exit(1);
    }

    let path = cfg_file.unwrap_or_else(|| src_path.join("config.toml"));
    let settings = format!(
        "# Includes one of the default files in src/bootstrap/defaults\n\
    profile = \"{}\"\n\
    changelog-seen = {}\n",
        profile, VERSION
    );
    t!(fs::write(path, settings));

    let include_path = profile.include_path(src_path);
    println!("`x.py` will now use the configuration at {}", include_path.display());

    let suggestions = match profile {
        Profile::Codegen | Profile::Compiler => &["check", "build", "test"][..],
        Profile::Library => &["check", "build", "test library/std", "doc"],
        Profile::User => &["dist", "build"],
    };

    println!();

    t!(install_git_hook_maybe(src_path));

    println!();

    println!("To get started, try one of the following commands:");
    for cmd in suggestions {
        println!("- `x.py {}`", cmd);
    }

    if profile != Profile::User {
        println!(
            "For more suggestions, see https://rustc-dev-guide.rust-lang.org/building/suggested.html"
        );
    }
}

// Used to get the path for `Subcommand::Setup`
pub fn interactive_path() -> io::Result<Profile> {
    fn abbrev_all() -> impl Iterator<Item = (String, Profile)> {
        ('a'..).map(|c| c.to_string()).zip(Profile::all())
    }

    fn parse_with_abbrev(input: &str) -> Result<Profile, String> {
        let input = input.trim().to_lowercase();
        for (letter, profile) in abbrev_all() {
            if input == letter {
                return Ok(profile);
            }
        }
        input.parse()
    }

    println!("Welcome to the Rust project! What do you want to do with x.py?");
    for (letter, profile) in abbrev_all() {
        println!("{}) {}: {}", letter, profile, profile.purpose());
    }
    let template = loop {
        print!(
            "Please choose one ({}): ",
            abbrev_all().map(|(l, _)| l).collect::<Vec<_>>().join("/")
        );
        io::stdout().flush()?;
        let mut input = String::new();
        io::stdin().read_line(&mut input)?;
        if input == "" {
            eprintln!("EOF on stdin, when expecting answer to question.  Giving up.");
            std::process::exit(1);
        }
        break match parse_with_abbrev(&input) {
            Ok(profile) => profile,
            Err(err) => {
                println!("error: {}", err);
                println!("note: press Ctrl+C to exit");
                continue;
            }
        };
    };
    Ok(template)
}

// install a git hook to automatically run tidy --bless, if they want
fn install_git_hook_maybe(src_path: &Path) -> io::Result<()> {
    let mut input = String::new();
    println!(
        "Rust's CI will automatically fail if it doesn't pass `tidy`, the internal tool for ensuring code quality.
If you'd like, x.py can install a git hook for you that will automatically run `tidy --bless` on each commit
to ensure your code is up to par. If you decide later that this behavior is undesirable,
simply delete the `pre-commit` file from .git/hooks."
    );

    let should_install = loop {
        print!("Would you like to install the git hook?: [y/N] ");
        io::stdout().flush()?;
        input.clear();
        io::stdin().read_line(&mut input)?;
        break match input.trim().to_lowercase().as_str() {
            "y" | "yes" => true,
            "n" | "no" | "" => false,
            _ => {
                println!("error: unrecognized option '{}'", input.trim());
                println!("note: press Ctrl+C to exit");
                continue;
            }
        };
    };

    Ok(if should_install {
        let src = src_path.join("src").join("etc").join("pre-commit.sh");
        let git = t!(Command::new("git").args(&["rev-parse", "--git-common-dir"]).output().map(
            |output| {
                assert!(output.status.success(), "failed to run `git`");
                PathBuf::from(t!(String::from_utf8(output.stdout)).trim())
            }
        ));
        let dst = git.join("hooks").join("pre-commit");
        match fs::hard_link(src, &dst) {
            Err(e) => println!(
                "error: could not create hook {}: do you already have the git hook installed?\n{}",
                dst.display(),
                e
            ),
            Ok(_) => println!("Linked `src/etc/pre-commit.sh` to `.git/hooks/pre-commit`"),
        };
    } else {
        println!("Ok, skipping installation!");
    })
}
<commit_msg>Create config.toml in the current directory, not the top-level directory<commit_after>use crate::{t, VERSION};
use std::fmt::Write as _;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::str::FromStr;
use std::{
    env, fmt, fs,
    io::{self, Write},
};

#[derive(Clone, Copy, Eq, PartialEq)]
pub enum Profile {
    Compiler,
    Codegen,
    Library,
    User,
}

impl Profile {
    fn include_path(&self, src_path: &Path) -> PathBuf {
        PathBuf::from(format!("{}/src/bootstrap/defaults/config.{}.toml", src_path.display(), self))
    }

    pub fn all() -> impl Iterator<Item = Self> {
        use Profile::*;
        // N.B. these are ordered by how they are displayed, not alphabetically
        [Library, Compiler, Codegen, User].iter().copied()
    }

    pub fn purpose(&self) -> String {
        use Profile::*;
        match self {
            Library => "Contribute to the standard library",
            Compiler => "Contribute to the compiler or rustdoc",
            Codegen => "Contribute to the compiler, and also modify LLVM or codegen",
            User => "Install Rust from source",
        }
        .to_string()
    }

    pub fn all_for_help(indent: &str) -> String {
        let mut out = String::new();
        for choice in Profile::all() {
            writeln!(&mut out, "{}{}: {}", indent, choice, choice.purpose()).unwrap();
        }
        out
    }
}

impl FromStr for Profile {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "lib" | "library" => Ok(Profile::Library),
            "compiler" | "rustdoc" => Ok(Profile::Compiler),
            "llvm" | "codegen" => Ok(Profile::Codegen),
            "maintainer" | "user" => Ok(Profile::User),
            _ => Err(format!("unknown profile: '{}'", s)),
        }
    }
}

impl fmt::Display for Profile {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Profile::Compiler => write!(f, "compiler"),
            Profile::Codegen => write!(f, "codegen"),
            Profile::Library => write!(f, "library"),
            Profile::User => write!(f, "user"),
        }
    }
}

pub fn setup(src_path: &Path, profile: Profile) {
    let cfg_file = env::var_os("BOOTSTRAP_CONFIG").map(PathBuf::from);

    if cfg_file.as_ref().map_or(false, |f| f.exists()) {
        let file = cfg_file.unwrap();
        println!(
            "error: you asked `x.py` to setup a new config file, but one already exists at `{}`",
            file.display()
        );
        println!("help: try adding `profile = \"{}\"` at the top of {}", profile, file.display());
        println!(
            "note: this will use the configuration in {}",
            profile.include_path(src_path).display()
        );
        std::process::exit(1);
    }

    let path = cfg_file.unwrap_or("config.toml".into());
    let settings = format!(
        "# Includes one of the default files in src/bootstrap/defaults\n\
    profile = \"{}\"\n\
    changelog-seen = {}\n",
        profile, VERSION
    );
    t!(fs::write(path, settings));

    let include_path = profile.include_path(src_path);
    println!("`x.py` will now use the configuration at {}", include_path.display());

    let suggestions = match profile {
        Profile::Codegen | Profile::Compiler => &["check", "build", "test"][..],
        Profile::Library => &["check", "build", "test library/std", "doc"],
        Profile::User => &["dist", "build"],
    };

    println!();

    t!(install_git_hook_maybe(src_path));

    println!();

    println!("To get started, try one of the following commands:");
    for cmd in suggestions {
        println!("- `x.py {}`", cmd);
    }

    if profile != Profile::User {
        println!(
            "For more suggestions, see https://rustc-dev-guide.rust-lang.org/building/suggested.html"
        );
    }
}

// Used to get the path for `Subcommand::Setup`
pub fn interactive_path() -> io::Result<Profile> {
    fn abbrev_all() -> impl Iterator<Item = (String, Profile)> {
        ('a'..).map(|c| c.to_string()).zip(Profile::all())
    }

    fn parse_with_abbrev(input: &str) -> Result<Profile, String> {
        let input = input.trim().to_lowercase();
        for (letter, profile) in abbrev_all() {
            if input == letter {
                return Ok(profile);
            }
        }
        input.parse()
    }

    println!("Welcome to the Rust project! What do you want to do with x.py?");
    for (letter, profile) in abbrev_all() {
        println!("{}) {}: {}", letter, profile, profile.purpose());
    }
    let template = loop {
        print!(
            "Please choose one ({}): ",
            abbrev_all().map(|(l, _)| l).collect::<Vec<_>>().join("/")
        );
        io::stdout().flush()?;
        let mut input = String::new();
        io::stdin().read_line(&mut input)?;
        if input == "" {
            eprintln!("EOF on stdin, when expecting answer to question.  Giving up.");
            std::process::exit(1);
        }
        break match parse_with_abbrev(&input) {
            Ok(profile) => profile,
            Err(err) => {
                println!("error: {}", err);
                println!("note: press Ctrl+C to exit");
                continue;
            }
        };
    };
    Ok(template)
}

// install a git hook to automatically run tidy --bless, if they want
fn install_git_hook_maybe(src_path: &Path) -> io::Result<()> {
    let mut input = String::new();
    println!(
        "Rust's CI will automatically fail if it doesn't pass `tidy`, the internal tool for ensuring code quality.
If you'd like, x.py can install a git hook for you that will automatically run `tidy --bless` on each commit
to ensure your code is up to par. If you decide later that this behavior is undesirable,
simply delete the `pre-commit` file from .git/hooks."
    );

    let should_install = loop {
        print!("Would you like to install the git hook?: [y/N] ");
        io::stdout().flush()?;
        input.clear();
        io::stdin().read_line(&mut input)?;
        break match input.trim().to_lowercase().as_str() {
            "y" | "yes" => true,
            "n" | "no" | "" => false,
            _ => {
                println!("error: unrecognized option '{}'", input.trim());
                println!("note: press Ctrl+C to exit");
                continue;
            }
        };
    };

    Ok(if should_install {
        let src = src_path.join("src").join("etc").join("pre-commit.sh");
        let git = t!(Command::new("git").args(&["rev-parse", "--git-common-dir"]).output().map(
            |output| {
                assert!(output.status.success(), "failed to run `git`");
                PathBuf::from(t!(String::from_utf8(output.stdout)).trim())
            }
        ));
        let dst = git.join("hooks").join("pre-commit");
        match fs::hard_link(src, &dst) {
            Err(e) => println!(
                "error: could not create hook {}: do you already have the git hook installed?\n{}",
                dst.display(),
                e
            ),
            Ok(_) => println!("Linked `src/etc/pre-commit.sh` to `.git/hooks/pre-commit`"),
        };
    } else {
        println!("Ok, skipping installation!");
    })
}
<|endoftext|>


<commit_before>extern crate glib;

extern crate gtk;
use gtk::{Inhibit, WidgetExt};

extern crate cairo;

use std::rc::{Rc, Weak};
use std::cell::RefCell;

use std::collections::vec_deque::{VecDeque};

use std::ops::{Deref, DerefMut};

use ::media::{Context, AudioBuffer};

use super::{MediaController, MediaHandler};

pub struct AudioController {
    media_ctl: MediaController,
    drawingarea: gtk::DrawingArea,

    circ_buffer: VecDeque<AudioBuffer>,
    sample_offset: f64,
    samples_nb: f64,
}

impl AudioController {
    pub fn new(builder: &gtk::Builder) -> Rc<RefCell<Self>> {
        let this = Rc::new(RefCell::new(AudioController {
            media_ctl: MediaController::new(
                builder.get_object("audio-container").unwrap(),
            ),
            drawingarea: builder.get_object("audio-drawingarea").unwrap(),

            circ_buffer: VecDeque::new(),
            sample_offset: 0f64,
            samples_nb: 0f64,
        }));

        {
            let this_ref = this.borrow();
            let this_weak = Rc::downgrade(&this);
            this_ref.drawingarea.connect_draw(move |ref drawing_area, ref cairo_ctx| {
                let this = this_weak.upgrade()
                    .expect("Main controller is no longer available for select_media");
                return this.borrow().draw(drawing_area, cairo_ctx);
            });
        }

        this
    }

    pub fn clear(&mut self) {
        self.circ_buffer.clear();
        self.sample_offset = 0f64;
        self.samples_nb = 0f64;
    }

    pub fn have_buffer(&mut self, buffer: AudioBuffer) {
        // Firt approximation: suppose the buffers come in ordered
        if self.sample_offset == 0f64 {
            self.sample_offset = buffer.sample_offset as f64;
        }
        self.samples_nb += buffer.samples_nb as f64;

        self.circ_buffer.push_back(buffer);
    }

    fn draw(&self, drawing_area: &gtk::DrawingArea, cr: &cairo::Context) -> Inhibit {
        if self.circ_buffer.len() == 0 {
            return Inhibit(false);
        }

        let sample_nb = self.samples_nb.min(5_000f64);
        let sample_dyn = 1024f64;

        let allocation = drawing_area.get_allocation();
        cr.scale(
            allocation.width as f64 / sample_nb,
            allocation.height as f64 / 2f64 / sample_dyn,
        );
        cr.set_line_width(1f64);

        for ref buffer in self.circ_buffer.iter() {
            for ref channel in buffer.channels.iter() {
                let colors = vec![(0.8f64, 0.8f64, 0.8f64), (0.8f64, 0f64, 0f64)][channel.id];
                cr.set_source_rgb(colors.0, colors.1, colors.2);

                let mut x = buffer.sample_offset as f64 - self.sample_offset;
                let mut is_first = true;
                for &sample in channel.iter() {
                    let y = sample_dyn * (1f64 - sample);
                    if !is_first {
                        cr.line_to(x, y);
                    } else {
                        cr.move_to(x, y);
                        is_first = false;
                    }
                    x += 1f64;
                    if x > sample_nb { break; }
                }

                cr.stroke();
            }
        }

        Inhibit(false)
    }
}

impl Deref for AudioController {
	type Target = MediaController;

	fn deref(&self) -> &Self::Target {
		&self.media_ctl
	}
}

impl DerefMut for AudioController {
	fn deref_mut(&mut self) -> &mut Self::Target {
		&mut self.media_ctl
	}
}

impl MediaHandler for AudioController {
    fn new_media(&mut self, ctx: &Context) {
        if let Some(audio_buffer) = self.circ_buffer.get(0) {
            println!("\n{:?}", audio_buffer.caps);
            self.drawingarea.queue_draw();
            self.media_ctl.show();
        }
        else {
            self.media_ctl.hide();
        }
    }
}
<commit_msg>Handle more than 2 audio channels (#7)<commit_after>extern crate glib;

extern crate gtk;
use gtk::{Inhibit, WidgetExt};

extern crate cairo;

use std::rc::{Rc, Weak};
use std::cell::RefCell;

use std::collections::vec_deque::{VecDeque};

use std::ops::{Deref, DerefMut};

use ::media::{Context, AudioBuffer};

use super::{MediaController, MediaHandler};

pub struct AudioController {
    media_ctl: MediaController,
    drawingarea: gtk::DrawingArea,

    circ_buffer: VecDeque<AudioBuffer>,
    channels: usize,
    sample_offset: f64,
    samples_nb: f64,
}

impl AudioController {
    pub fn new(builder: &gtk::Builder) -> Rc<RefCell<Self>> {
        let this = Rc::new(RefCell::new(AudioController {
            media_ctl: MediaController::new(
                builder.get_object("audio-container").unwrap(),
            ),
            drawingarea: builder.get_object("audio-drawingarea").unwrap(),

            circ_buffer: VecDeque::new(),
            channels: 0,
            sample_offset: 0f64,
            samples_nb: 0f64,
        }));

        {
            let this_ref = this.borrow();
            let this_weak = Rc::downgrade(&this);
            this_ref.drawingarea.connect_draw(move |ref drawing_area, ref cairo_ctx| {
                let this = this_weak.upgrade()
                    .expect("Main controller is no longer available for select_media");
                return this.borrow().draw(drawing_area, cairo_ctx);
            });
        }

        this
    }

    pub fn clear(&mut self) {
        self.circ_buffer.clear();
        self.channels = 0;
        self.sample_offset = 0f64;
        self.samples_nb = 0f64;
    }

    pub fn have_buffer(&mut self, buffer: AudioBuffer) {
        // Firt approximation: suppose the buffers come in ordered
        if self.sample_offset == 0f64 {
            self.sample_offset = buffer.sample_offset as f64;
            self.channels = buffer.caps.channels;
        }
        self.samples_nb += buffer.samples_nb as f64;

        self.circ_buffer.push_back(buffer);
    }

    fn draw(&self, drawing_area: &gtk::DrawingArea, cr: &cairo::Context) -> Inhibit {
        if self.circ_buffer.len() == 0 {
            return Inhibit(false);
        }

        let sample_nb = self.samples_nb.min(5_000f64);
        let sample_dyn = 1024f64;

        let allocation = drawing_area.get_allocation();
        cr.scale(
            allocation.width as f64 / sample_nb,
            allocation.height as f64 / 2f64 / sample_dyn,
        );
        cr.set_line_width(1f64);

        let mut colors = vec![(0.9f64, 0.9f64, 0.9f64), (0.9f64, 0f64, 0f64)];
        for channel in 2..self.channels {
            colors.push((0f64, 0f64, 0.2f64 * channel as f64));
        }

        for ref buffer in self.circ_buffer.iter() {
            for ref channel in buffer.channels.iter() {
                let color = colors[channel.id];
                cr.set_source_rgb(color.0, color.1, color.2);

                let mut x = buffer.sample_offset as f64 - self.sample_offset;
                let mut is_first = true;
                for &sample in channel.iter() {
                    let y = sample_dyn * (1f64 - sample);
                    if !is_first {
                        cr.line_to(x, y);
                    } else {
                        cr.move_to(x, y);
                        is_first = false;
                    }
                    x += 1f64;
                    if x > sample_nb { break; }
                }

                cr.stroke();
            }
        }

        Inhibit(false)
    }
}

impl Deref for AudioController {
	type Target = MediaController;

	fn deref(&self) -> &Self::Target {
		&self.media_ctl
	}
}

impl DerefMut for AudioController {
	fn deref_mut(&mut self) -> &mut Self::Target {
		&mut self.media_ctl
	}
}

impl MediaHandler for AudioController {
    fn new_media(&mut self, ctx: &Context) {
        if let Some(audio_buffer) = self.circ_buffer.get(0) {
            println!("\n{:?}", audio_buffer.caps);
            self.drawingarea.queue_draw();
            self.media_ctl.show();
        }
        else {
            self.media_ctl.hide();
        }
    }
}
<|endoftext|>


<commit_before>//! Module: czmq-zmsg

use {czmq_sys, Error, ErrorKind, RawInterface, Result, Sockish, ZFrame};
use std::{error, fmt, mem, ptr, result};
use std::ffi::{CStr, CString};

pub struct ZMsg {
    zmsg: *mut czmq_sys::zmsg_t,
    owned: bool,
}

unsafe impl Send for ZMsg {}
unsafe impl Sync for ZMsg {}

impl Drop for ZMsg {
    fn drop(&mut self) {
        if self.owned {
            unsafe { czmq_sys::zmsg_destroy(&mut self.zmsg) };
        }
    }
}

impl ZMsg {
    pub fn new() -> ZMsg {
        ZMsg {
            zmsg: unsafe { czmq_sys::zmsg_new() },
            owned: true,
        }
    }

    pub fn recv<S: Sockish>(source: &S) -> Result<ZMsg> {
        let zmsg = unsafe { czmq_sys::zmsg_recv(source.borrow_raw()) };

        if zmsg == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZMsg {
                zmsg: zmsg,
                owned: true,
            })
        }
    }

    // XXX We'll have to roll our own here as we can't imitate a C
    // file handle without C boilerplate...which we're not doing!
    // pub fn zmsg_load(file: *mut FILE) -> *mut zmsg_t;

    pub fn encode(&self) -> Result<ZFrame> {
        let zframe = unsafe { czmq_sys::zmsg_encode(self.zmsg) };

        if zframe == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZFrame::from_raw(zframe, true))
        }
    }

    pub fn decode(frame: &ZFrame) -> Result<ZMsg> {
        let zmsg = unsafe { czmq_sys::zmsg_decode(frame.borrow_raw()) };

        if zmsg == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZMsg {
                zmsg: zmsg,
                owned: true,
            })
        }
    }

    pub fn new_signal(status: u8) -> Result<ZMsg> {
        let zmsg = unsafe { czmq_sys::zmsg_new_signal(status) };

        if zmsg == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZMsg {
                zmsg: zmsg,
                owned: true,
            })
        }
    }

    pub fn send<D: Sockish>(self, dest: &D) -> Result<()> {
        let mut zmsg = self;

        let rc = unsafe { czmq_sys::zmsg_send(&mut zmsg.zmsg, dest.borrow_raw()) };
        if rc == -1 {
            Err(Error::new(ErrorKind::NonZero, ZMsgError::CmdFailed))
        } else {
            Ok(())
        }
    }

    // pub fn zmsg_sendm(self_p: *mut *mut zmsg_t,
    //                   dest: *mut ::std::os::raw::c_void)
    //  -> ::std::os::raw::c_int;

    pub fn size(&self) -> usize {
        unsafe { czmq_sys::zmsg_size(self.zmsg) as usize }
    }

    // pub fn zmsg_content_size(_self: *mut zmsg_t) -> size_t;
    // pub fn zmsg_prepend(_self: *mut zmsg_t, frame_p: *mut *mut zframe_t)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_append(_self: *mut zmsg_t, frame_p: *mut *mut zframe_t)
    //  -> ::std::os::raw::c_int;

    pub fn pop(&self) -> Option<ZFrame> {
        let ptr = unsafe { czmq_sys::zmsg_pop(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            Some(ZFrame::from_raw(ptr, true))
        }
    }

    // pub fn zmsg_pushmem(_self: *mut zmsg_t,
    //                     src: *const ::std::os::raw::c_void, size: size_t)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_addmem(_self: *mut zmsg_t, src: *const ::std::os::raw::c_void,
    //                    size: size_t) -> ::std::os::raw::c_int;
    // pub fn zmsg_pushstr(_self: *mut zmsg_t,
    //                     string: *const ::std::os::raw::c_char)
    //  -> ::std::os::raw::c_int;

    pub fn addstr(&self, string: &str) -> Result<()> {
        let string_c = CString::new(string).unwrap_or(CString::new("").unwrap());
        let rc = unsafe { czmq_sys::zmsg_addstr(self.zmsg, string_c.as_ptr()) };

        // Deliberately leak this memory, which will be managed by C
        mem::forget(string_c);

        if rc == -1 {
            Err(Error::new(ErrorKind::NonZero, ZMsgError::CmdFailed))
        } else {
            Ok(())
        }
    }

    // pub fn zmsg_pushstrf(_self: *mut zmsg_t,
    //                      format: *const ::std::os::raw::c_char, ...)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_addstrf(_self: *mut zmsg_t,
    //                     format: *const ::std::os::raw::c_char, ...)
    //  -> ::std::os::raw::c_int;

    pub fn popstr(&self) -> Option<result::Result<String, Vec<u8>>> {
        let ptr = unsafe { czmq_sys::zmsg_popstr(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            let c_string = unsafe { CStr::from_ptr(ptr).to_owned() };
            let bytes = c_string.as_bytes().to_vec();
            match c_string.into_string() {
                Ok(s) => Some(Ok(s)),
                Err(_) => Some(Err(bytes)),
            }
        }
    }

    pub fn popbytes(&self) -> Option<Vec<u8>> {
        let ptr = unsafe { czmq_sys::zmsg_popstr(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            let c_string = unsafe { CStr::from_ptr(ptr).to_owned() };
            Some(c_string.to_bytes().to_vec())
        }
    }

    // pub fn zmsg_addmsg(_self: *mut zmsg_t, msg_p: *mut *mut zmsg_t)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_popmsg(_self: *mut zmsg_t) -> *mut zmsg_t;
    // pub fn zmsg_remove(_self: *mut zmsg_t, frame: *mut zframe_t);
    // pub fn zmsg_first(_self: *mut zmsg_t) -> *mut zframe_t;
    // pub fn zmsg_next(_self: *mut zmsg_t) -> *mut zframe_t;
    // pub fn zmsg_last(_self: *mut zmsg_t) -> *mut zframe_t;
    // pub fn zmsg_save(_self: *mut zmsg_t, file: *mut FILE)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_dup(_self: *mut zmsg_t) -> *mut zmsg_t;
    // pub fn zmsg_print(_self: *mut zmsg_t);
    // pub fn zmsg_eq(_self: *mut zmsg_t, other: *mut zmsg_t) -> u8;

    pub fn signal(&self) -> Result<u8> {
        let signal = unsafe { czmq_sys::zmsg_signal(self.zmsg) };
        if signal == -1 {
            Err(Error::new(ErrorKind::NonZero, ZMsgError::CmdFailed))
        } else {
            Ok(signal as u8)
        }
    }

    // pub fn zmsg_test(verbose: u8);

    pub fn borrow_raw(&self) -> *mut czmq_sys::zmsg_t {
        self.zmsg
    }
}

impl RawInterface<czmq_sys::zmsg_t> for ZMsg {
    fn from_raw(ptr: *mut czmq_sys::zmsg_t, owned: bool) -> ZMsg {
        ZMsg {
            zmsg: ptr,
            owned: owned,
        }
    }

    fn into_raw(mut self) -> *mut czmq_sys::zmsg_t {
        self.owned = false;
        self.zmsg
    }

    fn borrow_raw(&self) -> *mut czmq_sys::zmsg_t {
        self.zmsg
    }
}

impl Iterator for ZMsg {
    type Item = ZFrame;

    fn next(&mut self) -> Option<Self::Item> {
        self.pop()
    }
}

#[derive(Debug)]
pub enum ZMsgError {
    CmdFailed,
}

impl fmt::Display for ZMsgError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ZMsgError::CmdFailed => write!(f, "ZMsg command failed"),
        }
    }
}

impl error::Error for ZMsgError {
    fn description(&self) -> &str {
        match *self {
            ZMsgError::CmdFailed => "ZMsg command failed",
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use {zmq, ZSock, zsys_init};

    #[test]
    fn test_sendrecv_zmq() {
        let mut ctx = zmq::Context::new();

        let mut server = ctx.socket(zmq::REP).unwrap();
        server.bind("inproc://zmsg_sendrecv_zmq").unwrap();

        let mut client = ctx.socket(zmq::REQ).unwrap();
        client.connect("inproc://zmsg_sendrecv_zmq").unwrap();

        let zmsg = ZMsg::new();
        zmsg.addstr("Hello world!").unwrap();
        zmsg.send(&client).unwrap();

        let zmsg_recv = ZMsg::recv(&server).unwrap();
        assert_eq!(zmsg_recv.popstr().unwrap().unwrap(), "Hello world!");
    }

    #[test]
    fn test_sendrecv_zsock() {
        zsys_init();

        let server = ZSock::new_rep("inproc://zmsg_sendrecv_zsock").unwrap();
        let client = ZSock::new_req("inproc://zmsg_sendrecv_zsock").unwrap();

        let zmsg = ZMsg::new();
        zmsg.addstr("Hello world!").unwrap();
        zmsg.send(&client).unwrap();

        let msg = ZMsg::recv(&server).unwrap();
        assert_eq!(msg.popstr().unwrap().unwrap(), "Hello world!");
    }

    #[test]
    fn test_encode_decode() {
        let msg = ZMsg::new();
        msg.addstr("moo").unwrap();
        let msg_decoded = ZMsg::decode(&msg.encode().unwrap()).unwrap();
        assert_eq!(msg_decoded.popstr().unwrap().unwrap(), "moo");
    }

    #[test]
    fn test_signal() {
        let msg = ZMsg::new_signal(97).unwrap();
        assert_eq!(msg.signal().unwrap(), 97);
    }

    #[test]
    fn test_size() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.size(), 1);
    }

    #[test]
    fn test_pop() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.pop().unwrap().data().unwrap().unwrap(), "123");
    }

    #[test]
    fn test_popstr() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.popstr().unwrap().unwrap(), "123");
    }

    #[test]
    fn test_popbytes() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.popbytes().unwrap(), "123".as_bytes());
    }
}
<commit_msg>Implemented iter fns for ZMsg so that Iterator doesn't pop frames from msg<commit_after>//! Module: czmq-zmsg

use {czmq_sys, Error, ErrorKind, RawInterface, Result, Sockish, ZFrame};
use std::{error, fmt, mem, ptr, result};
use std::ffi::{CStr, CString};

pub struct ZMsg {
    zmsg: *mut czmq_sys::zmsg_t,
    owned: bool,
}

unsafe impl Send for ZMsg {}
unsafe impl Sync for ZMsg {}

impl Drop for ZMsg {
    fn drop(&mut self) {
        if self.owned {
            unsafe { czmq_sys::zmsg_destroy(&mut self.zmsg) };
        }
    }
}

impl ZMsg {
    pub fn new() -> ZMsg {
        ZMsg {
            zmsg: unsafe { czmq_sys::zmsg_new() },
            owned: true,
        }
    }

    pub fn recv<S: Sockish>(source: &S) -> Result<ZMsg> {
        let zmsg = unsafe { czmq_sys::zmsg_recv(source.borrow_raw()) };

        if zmsg == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZMsg {
                zmsg: zmsg,
                owned: true,
            })
        }
    }

    // XXX We'll have to roll our own here as we can't imitate a C
    // file handle without C boilerplate...which we're not doing!
    // pub fn zmsg_load(file: *mut FILE) -> *mut zmsg_t;

    pub fn encode(&self) -> Result<ZFrame> {
        let zframe = unsafe { czmq_sys::zmsg_encode(self.zmsg) };

        if zframe == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZFrame::from_raw(zframe, true))
        }
    }

    pub fn decode(frame: &ZFrame) -> Result<ZMsg> {
        let zmsg = unsafe { czmq_sys::zmsg_decode(frame.borrow_raw()) };

        if zmsg == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZMsg {
                zmsg: zmsg,
                owned: true,
            })
        }
    }

    pub fn new_signal(status: u8) -> Result<ZMsg> {
        let zmsg = unsafe { czmq_sys::zmsg_new_signal(status) };

        if zmsg == ptr::null_mut() {
            Err(Error::new(ErrorKind::NullPtr, ZMsgError::CmdFailed))
        } else {
            Ok(ZMsg {
                zmsg: zmsg,
                owned: true,
            })
        }
    }

    pub fn send<D: Sockish>(self, dest: &D) -> Result<()> {
        let mut zmsg = self;

        let rc = unsafe { czmq_sys::zmsg_send(&mut zmsg.zmsg, dest.borrow_raw()) };
        if rc == -1 {
            Err(Error::new(ErrorKind::NonZero, ZMsgError::CmdFailed))
        } else {
            Ok(())
        }
    }

    // pub fn zmsg_sendm(self_p: *mut *mut zmsg_t,
    //                   dest: *mut ::std::os::raw::c_void)
    //  -> ::std::os::raw::c_int;

    pub fn size(&self) -> usize {
        unsafe { czmq_sys::zmsg_size(self.zmsg) as usize }
    }

    // pub fn zmsg_content_size(_self: *mut zmsg_t) -> size_t;
    // pub fn zmsg_prepend(_self: *mut zmsg_t, frame_p: *mut *mut zframe_t)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_append(_self: *mut zmsg_t, frame_p: *mut *mut zframe_t)
    //  -> ::std::os::raw::c_int;

    pub fn pop(&self) -> Option<ZFrame> {
        let ptr = unsafe { czmq_sys::zmsg_pop(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            Some(ZFrame::from_raw(ptr, true))
        }
    }

    // pub fn zmsg_pushmem(_self: *mut zmsg_t,
    //                     src: *const ::std::os::raw::c_void, size: size_t)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_addmem(_self: *mut zmsg_t, src: *const ::std::os::raw::c_void,
    //                    size: size_t) -> ::std::os::raw::c_int;
    // pub fn zmsg_pushstr(_self: *mut zmsg_t,
    //                     string: *const ::std::os::raw::c_char)
    //  -> ::std::os::raw::c_int;

    pub fn addstr(&self, string: &str) -> Result<()> {
        let string_c = CString::new(string).unwrap_or(CString::new("").unwrap());
        let rc = unsafe { czmq_sys::zmsg_addstr(self.zmsg, string_c.as_ptr()) };

        // Deliberately leak this memory, which will be managed by C
        mem::forget(string_c);

        if rc == -1 {
            Err(Error::new(ErrorKind::NonZero, ZMsgError::CmdFailed))
        } else {
            Ok(())
        }
    }

    // pub fn zmsg_pushstrf(_self: *mut zmsg_t,
    //                      format: *const ::std::os::raw::c_char, ...)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_addstrf(_self: *mut zmsg_t,
    //                     format: *const ::std::os::raw::c_char, ...)
    //  -> ::std::os::raw::c_int;

    pub fn popstr(&self) -> Option<result::Result<String, Vec<u8>>> {
        let ptr = unsafe { czmq_sys::zmsg_popstr(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            let c_string = unsafe { CStr::from_ptr(ptr).to_owned() };
            let bytes = c_string.as_bytes().to_vec();
            match c_string.into_string() {
                Ok(s) => Some(Ok(s)),
                Err(_) => Some(Err(bytes)),
            }
        }
    }

    pub fn popbytes(&self) -> Option<Vec<u8>> {
        let ptr = unsafe { czmq_sys::zmsg_popstr(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            let c_string = unsafe { CStr::from_ptr(ptr).to_owned() };
            Some(c_string.to_bytes().to_vec())
        }
    }

    // pub fn zmsg_addmsg(_self: *mut zmsg_t, msg_p: *mut *mut zmsg_t)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_popmsg(_self: *mut zmsg_t) -> *mut zmsg_t;
    // pub fn zmsg_remove(_self: *mut zmsg_t, frame: *mut zframe_t);

    pub fn first(&self) -> Option<ZFrame> {
        let ptr = unsafe { czmq_sys::zmsg_first(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            Some(ZFrame::from_raw(ptr, false))
        }
    }

    pub fn next(&self) -> Option<ZFrame> {
        let ptr = unsafe { czmq_sys::zmsg_next(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            Some(ZFrame::from_raw(ptr, false))
        }
    }

    // We can't call this fn last() as it conflicts with
    // Iterator::last(), which is implemented for ZMsg.
    pub fn ref_last(&self) -> Option<ZFrame> {
        let ptr = unsafe { czmq_sys::zmsg_last(self.zmsg) };

        if ptr == ptr::null_mut() {
            None
        } else {
            Some(ZFrame::from_raw(ptr, false))
        }
    }

    // pub fn zmsg_save(_self: *mut zmsg_t, file: *mut FILE)
    //  -> ::std::os::raw::c_int;
    // pub fn zmsg_dup(_self: *mut zmsg_t) -> *mut zmsg_t;
    // pub fn zmsg_print(_self: *mut zmsg_t);
    // pub fn zmsg_eq(_self: *mut zmsg_t, other: *mut zmsg_t) -> u8;

    pub fn signal(&self) -> Result<u8> {
        let signal = unsafe { czmq_sys::zmsg_signal(self.zmsg) };
        if signal == -1 {
            Err(Error::new(ErrorKind::NonZero, ZMsgError::CmdFailed))
        } else {
            Ok(signal as u8)
        }
    }

    // pub fn zmsg_test(verbose: u8);

    pub fn borrow_raw(&self) -> *mut czmq_sys::zmsg_t {
        self.zmsg
    }
}

impl RawInterface<czmq_sys::zmsg_t> for ZMsg {
    fn from_raw(ptr: *mut czmq_sys::zmsg_t, owned: bool) -> ZMsg {
        ZMsg {
            zmsg: ptr,
            owned: owned,
        }
    }

    fn into_raw(mut self) -> *mut czmq_sys::zmsg_t {
        self.owned = false;
        self.zmsg
    }

    fn borrow_raw(&self) -> *mut czmq_sys::zmsg_t {
        self.zmsg
    }
}

impl Iterator for ZMsg {
    type Item = ZFrame;

    fn next(&mut self) -> Option<Self::Item> {
        ZMsg::next(self)
    }
}

#[derive(Debug)]
pub enum ZMsgError {
    CmdFailed,
}

impl fmt::Display for ZMsgError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ZMsgError::CmdFailed => write!(f, "ZMsg command failed"),
        }
    }
}

impl error::Error for ZMsgError {
    fn description(&self) -> &str {
        match *self {
            ZMsgError::CmdFailed => "ZMsg command failed",
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use {zmq, ZSock, zsys_init};

    #[test]
    fn test_sendrecv_zmq() {
        let mut ctx = zmq::Context::new();

        let mut server = ctx.socket(zmq::REP).unwrap();
        server.bind("inproc://zmsg_sendrecv_zmq").unwrap();

        let mut client = ctx.socket(zmq::REQ).unwrap();
        client.connect("inproc://zmsg_sendrecv_zmq").unwrap();

        let zmsg = ZMsg::new();
        zmsg.addstr("Hello world!").unwrap();
        zmsg.send(&client).unwrap();

        let zmsg_recv = ZMsg::recv(&server).unwrap();
        assert_eq!(zmsg_recv.popstr().unwrap().unwrap(), "Hello world!");
    }

    #[test]
    fn test_sendrecv_zsock() {
        zsys_init();

        let server = ZSock::new_rep("inproc://zmsg_sendrecv_zsock").unwrap();
        let client = ZSock::new_req("inproc://zmsg_sendrecv_zsock").unwrap();

        let zmsg = ZMsg::new();
        zmsg.addstr("Hello world!").unwrap();
        zmsg.send(&client).unwrap();

        let msg = ZMsg::recv(&server).unwrap();
        assert_eq!(msg.popstr().unwrap().unwrap(), "Hello world!");
    }

    #[test]
    fn test_encode_decode() {
        let msg = ZMsg::new();
        msg.addstr("moo").unwrap();
        let msg_decoded = ZMsg::decode(&msg.encode().unwrap()).unwrap();
        assert_eq!(msg_decoded.popstr().unwrap().unwrap(), "moo");
    }

    #[test]
    fn test_signal() {
        let msg = ZMsg::new_signal(97).unwrap();
        assert_eq!(msg.signal().unwrap(), 97);
    }

    #[test]
    fn test_size() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.size(), 1);
    }

    #[test]
    fn test_pop() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.pop().unwrap().data().unwrap().unwrap(), "123");
    }

    #[test]
    fn test_popstr() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.popstr().unwrap().unwrap(), "123");
    }

    #[test]
    fn test_popbytes() {
        let msg = ZMsg::new();
        msg.addstr("123").unwrap();
        assert_eq!(msg.popbytes().unwrap(), "123".as_bytes());
    }

    #[test]
    fn test_iter_fns() {
        let msg = ZMsg::new();
        msg.addstr("1").unwrap();
        msg.addstr("2").unwrap();
        msg.addstr("3").unwrap();

        let frame = msg.first().unwrap();
        assert_eq!(frame.data().unwrap().unwrap(), "1");

        let frame = msg.next().unwrap();
        assert_eq!(frame.data().unwrap().unwrap(), "2");

        let frame = msg.ref_last().unwrap();
        assert_eq!(frame.data().unwrap().unwrap(), "3");
    }

    #[test]
    fn test_iter() {
        let msg = ZMsg::new();
        msg.addstr("1").unwrap();
        msg.addstr("2").unwrap();
        msg.addstr("3").unwrap();

        let mut i = 1;
        for x in msg {
            assert_eq!(x.data().unwrap().unwrap(), i.to_string());
            i += 1;
        }
    }
}
<|endoftext|>