<commit_before>use nom::{
  Consumer, ConsumerState,
  ErrorKind,
  HexDisplay,
  Input, IResult,
  Move, Needed,
};

use metadata;

use metadata::{Metadata, StreamInfo, metadata_parser};
use frame::{frame_parser, Frame};

enum ParserState {
  Marker,
  Metadata,
  Frame,
}

pub struct Stream {
  info: StreamInfo,
  pub metadata: Vec<Metadata>,
  pub frames: Vec<Frame>,
  state: ParserState,
  consumer_state: ConsumerState<(), ErrorKind, Move>,
}

named!(pub stream_parser <&[u8], Stream>,
  chain!(
    blocks: metadata_parser ~
    frames: many1!(apply!(frame_parser, &blocks.0)),
    move|| {
      let consumed = Move::Consume(0);

      Stream {
        info: blocks.0,
        metadata: blocks.1,
        frames: frames,
        state: ParserState::Marker,
        consumer_state: ConsumerState::Continue(consumed),
      }
    }
  )
);

impl Stream {
  pub fn new() -> Stream {
    let consumed = Move::Consume(0);

    Stream {
      info: StreamInfo::new(),
      metadata: Vec::new(),
      frames: Vec::new(),
      state: ParserState::Marker,
      consumer_state: ConsumerState::Continue(consumed),
    }
  }

  pub fn info(&self) -> StreamInfo {
    self.info
  }

  fn handle_marker(&mut self, input: &[u8]) {
    match tag!(input, "fLaC") {
      IResult::Done(i, _)       => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);

        self.state          = ParserState::Metadata;
        self.consumer_state = ConsumerState::Continue(consumed);
      }
      IResult::Error(_)         => {
        let kind = ErrorKind::Custom(0);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(size) => {
        let needed = Move::Await(size);

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }

  fn handle_metadata(&mut self, input: &[u8]) {
    match metadata::block(input) {
      IResult::Done(i, block) => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);
        let is_last  = block.is_last;

        if let metadata::Data::StreamInfo(info) = block.data {
          self.info = info;
        } else {
          self.metadata.push(block);
        }

        if is_last {
          self.state = ParserState::Frame;
        }

        self.consumer_state = ConsumerState::Continue(consumed);
      }
      IResult::Error(_)       => {
        let kind = ErrorKind::Custom(1);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(s)  => {
        let size = if let Needed::Size(length) = s {
          length
        } else {
          1024
        };
        let needed = Move::Await(Needed::Size(size));

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }

  fn handle_frame(&mut self, input: &[u8]) {
    match frame_parser(input, &self.info) {
      IResult::Done(i, frame) => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);

        self.frames.push(frame);

        self.consumer_state = ConsumerState::Continue(consumed);
      }
      IResult::Error(_)       => {
        let kind = ErrorKind::Custom(2);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(s)  => {
        let size = if let Needed::Size(length) = s {
          length
        } else {
          self.info.max_frame_size as usize
        };
        let needed = Move::Await(Needed::Size(size));

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }
}

impl<'a> Consumer<&'a [u8], (), ErrorKind, Move> for Stream {
  fn state(&self) -> &ConsumerState<(), ErrorKind, Move> {
    &self.consumer_state
  }

  fn handle(&mut self, input: Input<&'a [u8]>)
            -> &ConsumerState<(), ErrorKind, Move> {
    match input {
      Input::Element(i) | Input::Eof(Some(i)) => {
        match self.state {
          ParserState::Marker   => self.handle_marker(i),
          ParserState::Metadata => self.handle_metadata(i),
          ParserState::Frame    => self.handle_frame(i),
        }
      }
      Input::Empty | Input::Eof(None)         => {
        self.consumer_state = match self.state {
          ParserState::Marker   => ConsumerState::Error(ErrorKind::Custom(0)),
          ParserState::Metadata => ConsumerState::Error(ErrorKind::Custom(1)),
          ParserState::Frame    => ConsumerState::Done(Move::Consume(0), ()),
        };
      }
    }

    &self.consumer_state
  }
}
<commit_msg>Add full parse of a file with `Stream::from_file`<commit_after>use nom::{
  Consumer, ConsumerState,
  ErrorKind,
  HexDisplay,
  Producer, FileProducer,
  Input, IResult,
  Move, Needed,
};

use metadata;

use metadata::{Metadata, StreamInfo, metadata_parser};
use frame::{frame_parser, Frame};
use utility::resize_producer;

use std::io;
use std::io::{Error, Result};

enum ParserState {
  Marker,
  Metadata,
  Frame,
}

pub struct Stream {
  info: StreamInfo,
  pub metadata: Vec<Metadata>,
  pub frames: Vec<Frame>,
  state: ParserState,
  consumer_state: ConsumerState<(), ErrorKind, Move>,
}

named!(pub stream_parser <&[u8], Stream>,
  chain!(
    blocks: metadata_parser ~
    frames: many1!(apply!(frame_parser, &blocks.0)),
    move|| {
      let consumed = Move::Consume(0);

      Stream {
        info: blocks.0,
        metadata: blocks.1,
        frames: frames,
        state: ParserState::Marker,
        consumer_state: ConsumerState::Continue(consumed),
      }
    }
  )
);

impl Stream {
  pub fn new() -> Stream {
    let consumed = Move::Consume(0);

    Stream {
      info: StreamInfo::new(),
      metadata: Vec::new(),
      frames: Vec::new(),
      state: ParserState::Marker,
      consumer_state: ConsumerState::Continue(consumed),
    }
  }

  pub fn info(&self) -> StreamInfo {
    self.info
  }

  pub fn from_file(filename: &str) -> Result<Stream> {
    FileProducer::new(filename, 1024).and_then(|mut producer| {
      let consumed        = Move::Consume(0);
      let mut buffer_size = 1024;
      let mut is_error    = false;
      let mut stream      = Stream {
        info: StreamInfo::new(),
        metadata: Vec::new(),
        frames: Vec::new(),
        state: ParserState::Marker,
        consumer_state: ConsumerState::Continue(consumed),
      };

      loop {
        match *producer.apply(&mut stream) {
          ConsumerState::Done(_, _)      => break,
          ConsumerState::Continue(await) => {
            let result = resize_producer(&mut producer, &await, buffer_size);

            if let Some(size) = result {
              buffer_size = size;
            }

            continue;
          }
          ConsumerState::Error(_)        => {
            is_error = true;

            break;
          }
        }
      }

      if !is_error {
        Ok(stream)
      } else {
        let error_str = format!("parser: couldn't parse the given file {}",
                                filename);

        Err(Error::new(io::ErrorKind::InvalidData, error_str))
      }
    })
  }

  fn handle_marker(&mut self, input: &[u8]) {
    match tag!(input, "fLaC") {
      IResult::Done(i, _)       => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);

        self.state          = ParserState::Metadata;
        self.consumer_state = ConsumerState::Continue(consumed);
      }
      IResult::Error(_)         => {
        let kind = ErrorKind::Custom(0);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(size) => {
        let needed = Move::Await(size);

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }

  fn handle_metadata(&mut self, input: &[u8]) {
    match metadata::block(input) {
      IResult::Done(i, block) => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);
        let is_last  = block.is_last;

        if let metadata::Data::StreamInfo(info) = block.data {
          self.info = info;
        } else {
          self.metadata.push(block);
        }

        if is_last {
          self.state = ParserState::Frame;
        }

        self.consumer_state = ConsumerState::Continue(consumed);
      }
      IResult::Error(_)       => {
        let kind = ErrorKind::Custom(1);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(s)  => {
        let size = if let Needed::Size(length) = s {
          length
        } else {
          1024
        };
        let needed = Move::Await(Needed::Size(size));

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }

  fn handle_frame(&mut self, input: &[u8]) {
    match frame_parser(input, &self.info) {
      IResult::Done(i, frame) => {
        let offset   = input.offset(i);
        let consumed = Move::Consume(offset);

        self.frames.push(frame);

        self.consumer_state = ConsumerState::Continue(consumed);
      }
      IResult::Error(_)       => {
        let kind = ErrorKind::Custom(2);

        self.consumer_state = ConsumerState::Error(kind);
      }
      IResult::Incomplete(s)  => {
        let size = if let Needed::Size(length) = s {
          length
        } else {
          self.info.max_frame_size as usize
        };
        let needed = Move::Await(Needed::Size(size));

        self.consumer_state = ConsumerState::Continue(needed);
      }
    }
  }
}

impl<'a> Consumer<&'a [u8], (), ErrorKind, Move> for Stream {
  fn state(&self) -> &ConsumerState<(), ErrorKind, Move> {
    &self.consumer_state
  }

  fn handle(&mut self, input: Input<&'a [u8]>)
            -> &ConsumerState<(), ErrorKind, Move> {
    match input {
      Input::Element(i) | Input::Eof(Some(i)) => {
        match self.state {
          ParserState::Marker   => self.handle_marker(i),
          ParserState::Metadata => self.handle_metadata(i),
          ParserState::Frame    => self.handle_frame(i),
        }
      }
      Input::Empty | Input::Eof(None)         => {
        self.consumer_state = match self.state {
          ParserState::Marker   => ConsumerState::Error(ErrorKind::Custom(0)),
          ParserState::Metadata => ConsumerState::Error(ErrorKind::Custom(1)),
          ParserState::Frame    => ConsumerState::Done(Move::Consume(0), ()),
        };
      }
    }

    &self.consumer_state
  }
}
<|endoftext|>


<commit_before>//! Group membership API.
use rdsys;
use rdsys::types::*;

use std::ffi::CStr;
use std::slice;
use std::fmt;

/// Group member information container.
pub struct GroupMemberInfo(RDKafkaGroupMemberInfo);

impl GroupMemberInfo {
    /// Return the id of the member.
    pub fn id(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.member_id).to_str()
                .expect("Member id is not a valid UTF-8 string")
        }
    }

    /// Return the client id of the member.
    pub fn client_id(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.client_id).to_str()
                .expect("Client id is not a valid UTF-8 string")
        }
    }

    /// Return the client host of the member.
    pub fn client_host(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.client_host).to_str()
                .expect("Member host is not a valid UTF-8 string")
        }
    }
}

/// Group information container.
pub struct GroupInfo(RDKafkaGroupInfo);

impl GroupInfo {
    /// Return the name of the group.
    pub fn name(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.group).to_str()
                .expect("Group name is not a valid UTF-8 string")
        }
    }

    /// Returns the members of the group.
    pub fn members(&self) -> &[GroupMemberInfo] {
        unsafe { slice::from_raw_parts(self.0.members as *const GroupMemberInfo, self.0.member_cnt as usize) }
    }

    /// Returns the state of the group.
    pub fn state(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.state).to_str()
                .expect("State is not a valid UTF-8 string")
        }
    }

    /// Returns the protocol of the group.
    pub fn protocol(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.protocol).to_str()
                .expect("Protocol name is not a valid UTF-8 string")
        }
    }

    /// Returns the protocol type of the group.
    pub fn protocol_type(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.protocol_type).to_str()
                .expect("Protocol type is not a valid UTF-8 string")
        }
    }
}

impl fmt::Debug for GroupInfo {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.name())
    }
}

/// List of groups. This structure wraps the  pointer returned by rdkafka-sys, and deallocates all
/// the native resources when dropped.
pub struct GroupList(*const RDKafkaGroupList);

impl GroupList {
    /// Creates a new group list given a pointer to the native rdkafka-sys group list structure.
    pub fn from_ptr(ptr: *const RDKafkaGroupList) -> GroupList {
        GroupList(ptr)
    }

    /// Returns all the groups in the list.
    pub fn groups(&self) -> &[GroupInfo] {
        unsafe { slice::from_raw_parts((*self.0).groups as *const GroupInfo, (*self.0).group_cnt as usize) }
    }
}

impl Drop for GroupList {
    fn drop(&mut self) {
        unsafe { rdsys::rd_kafka_group_list_destroy(self.0) };
    }
}
<commit_msg>Expose metadata and assignment on GroupMemberInfo<commit_after>//! Group membership API.
use rdsys;
use rdsys::types::*;

use std::ffi::CStr;
use std::slice;
use std::fmt;

/// Group member information container.
pub struct GroupMemberInfo(RDKafkaGroupMemberInfo);

impl GroupMemberInfo {
    /// Return the id of the member.
    pub fn id(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.member_id).to_str()
                .expect("Member id is not a valid UTF-8 string")
        }
    }

    /// Return the client id of the member.
    pub fn client_id(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.client_id).to_str()
                .expect("Client id is not a valid UTF-8 string")
        }
    }

    /// Return the client host of the member.
    pub fn client_host(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.client_host).to_str()
                .expect("Member host is not a valid UTF-8 string")
        }
    }

    /// Return the metadata of the member
    pub fn metadata(&self) -> Option<&[u8]> {
        unsafe {
            if self.0.member_metadata.is_null() {
                None
            } else {
                Some(slice::from_raw_parts::<u8>(self.0.member_metadata as *const u8, self.0.member_metadata_size as usize))
            }
        }
    }

    /// Return the assignment of the member
    pub fn assignment(&self) -> Option<&[u8]> {
        unsafe {
            if self.0.member_assignment.is_null() {
                None
            } else {
                Some(slice::from_raw_parts::<u8>(self.0.member_assignment as *const u8, self.0.member_assignment_size as usize))
            }
        }
    }
}

/// Group information container.
pub struct GroupInfo(RDKafkaGroupInfo);

impl GroupInfo {
    /// Return the name of the group.
    pub fn name(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.group).to_str()
                .expect("Group name is not a valid UTF-8 string")
        }
    }

    /// Returns the members of the group.
    pub fn members(&self) -> &[GroupMemberInfo] {
        unsafe { slice::from_raw_parts(self.0.members as *const GroupMemberInfo, self.0.member_cnt as usize) }
    }

    /// Returns the state of the group.
    pub fn state(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.state).to_str()
                .expect("State is not a valid UTF-8 string")
        }
    }

    /// Returns the protocol of the group.
    pub fn protocol(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.protocol).to_str()
                .expect("Protocol name is not a valid UTF-8 string")
        }
    }

    /// Returns the protocol type of the group.
    pub fn protocol_type(&self) -> &str {
        unsafe {
            CStr::from_ptr(self.0.protocol_type).to_str()
                .expect("Protocol type is not a valid UTF-8 string")
        }
    }
}

impl fmt::Debug for GroupInfo {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.name())
    }
}

/// List of groups. This structure wraps the  pointer returned by rdkafka-sys, and deallocates all
/// the native resources when dropped.
pub struct GroupList(*const RDKafkaGroupList);

impl GroupList {
    /// Creates a new group list given a pointer to the native rdkafka-sys group list structure.
    pub fn from_ptr(ptr: *const RDKafkaGroupList) -> GroupList {
        GroupList(ptr)
    }

    /// Returns all the groups in the list.
    pub fn groups(&self) -> &[GroupInfo] {
        unsafe { slice::from_raw_parts((*self.0).groups as *const GroupInfo, (*self.0).group_cnt as usize) }
    }
}

impl Drop for GroupList {
    fn drop(&mut self) {
        unsafe { rdsys::rd_kafka_group_list_destroy(self.0) };
    }
}
<|endoftext|>


<commit_before>use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::collections::{BinaryHeap, HashMap};

use std::hash::Hash;

use super::visit::{EdgeRef, GraphBase, IntoEdges};
use crate::scored::MinScored;

use crate::algo::Measure;

/// \[Generic\] A* shortest path algorithm.
///
/// Computes the shortest path from `start` to `finish`, including the total path cost.
///
/// `finish` is implicitly given via the `is_goal` callback, which should return `true` if the
/// given node is the finish node.
///
/// The function `edge_cost` should return the cost for a particular edge. Edge costs must be
/// non-negative.
///
/// The function `estimate_cost` should return the estimated cost to the finish for a particular
/// node. For the algorithm to find the actual shortest path, it should be admissible, meaning that
/// it should never overestimate the actual cost to get to the nearest goal node. Estimate costs
/// must also be non-negative.
///
/// The graph should be `Visitable` and implement `IntoEdges`.
///
/// # Example
/// ```
/// use petgraph::Graph;
/// use petgraph::algo::astar;
///
/// let mut g = Graph::new();
/// let a = g.add_node((0., 0.));
/// let b = g.add_node((2., 0.));
/// let c = g.add_node((1., 1.));
/// let d = g.add_node((0., 2.));
/// let e = g.add_node((3., 3.));
/// let f = g.add_node((4., 2.));
/// g.extend_with_edges(&[
///     (a, b, 2),
///     (a, d, 4),
///     (b, c, 1),
///     (b, f, 7),
///     (c, e, 5),
///     (e, f, 1),
///     (d, e, 1),
/// ]);
///
/// // Graph represented with the weight of each edge
/// // Edges with '*' are part of the optimal path.
/// //
/// //     2       1
/// // a ----- b ----- c
/// // | 4*    | 7     |
/// // d       f       | 5
/// // | 1*    | 1*    |
/// // \------ e ------/
///
/// let path = astar(&g, a, |finish| finish == f, |e| *e.weight(), |_| 0);
/// assert_eq!(path, Some((6, vec![a, d, e, f])));
/// ```
///
/// Returns the total cost + the path of subsequent `NodeId` from start to finish, if one was
/// found.
pub fn astar<G, F, H, K, IsGoal>(
    graph: G,
    start: G::NodeId,
    mut is_goal: IsGoal,
    mut edge_cost: F,
    mut estimate_cost: H,
) -> Option<(K, Vec<G::NodeId>)>
where
    G: IntoEdges,
    IsGoal: FnMut(G::NodeId) -> bool,
    G::NodeId: Eq + Hash,
    F: FnMut(G::EdgeRef) -> K,
    H: FnMut(G::NodeId) -> K,
    K: Measure + Copy,
{
    let mut visit_next = BinaryHeap::new();
    let mut scores = HashMap::new(); // g-values, cost to reach the node
    let mut estimate_scores = HashMap::new(); // f-values, cost to reach + estimate cost to goal
    let mut path_tracker = PathTracker::<G>::new();

    let zero_score = K::default();
    scores.insert(start, zero_score);
    visit_next.push(MinScored(estimate_cost(start), start));

    while let Some(MinScored(estimate_score, node)) = visit_next.pop() {
        if is_goal(node) {
            let path = path_tracker.reconstruct_path_to(node);
            let cost = scores[&node];
            return Some((cost, path));
        }

        // This lookup can be unwrapped without fear of panic since the node was necessarily scored
        // before adding it to `visit_next`.
        let node_score = scores[&node];

        match estimate_scores.entry(node) {
            Occupied(mut entry) => {
                // If the node has already been visited with an equal or lower score than now, then
                // we do not need to re-visit it.
                if *entry.get() <= estimate_score {
                    continue;
                }
                entry.insert(estimate_score);
            }
            Vacant(entry) => {
                entry.insert(estimate_score);
            }
        }

        for edge in graph.edges(node) {
            let next = edge.target();
            let next_score = node_score + edge_cost(edge);

            match scores.entry(next) {
                Occupied(mut entry) => {
                    // No need to add neighbors that we have already reached through a shorter path
                    // than now.
                    if *entry.get() <= next_score {
                        continue;
                    }
                    entry.insert(next_score);
                }
                Vacant(entry) => {
                    entry.insert(next_score);
                }
            }

            path_tracker.set_predecessor(next, node);
            let next_estimate_score = next_score + estimate_cost(next);
            visit_next.push(MinScored(next_estimate_score, next));
        }
    }

    None
}

struct PathTracker<G>
where
    G: GraphBase,
    G::NodeId: Eq + Hash,
{
    came_from: HashMap<G::NodeId, G::NodeId>,
}

impl<G> PathTracker<G>
where
    G: GraphBase,
    G::NodeId: Eq + Hash,
{
    fn new() -> PathTracker<G> {
        PathTracker {
            came_from: HashMap::new(),
        }
    }

    fn set_predecessor(&mut self, node: G::NodeId, previous: G::NodeId) {
        self.came_from.insert(node, previous);
    }

    fn reconstruct_path_to(&self, last: G::NodeId) -> Vec<G::NodeId> {
        let mut path = vec![last];

        let mut current = last;
        while let Some(&previous) = self.came_from.get(&current) {
            path.push(previous);
            current = previous;
        }

        path.reverse();

        path
    }
}
<commit_msg>Apply suggestions - keep Visitable on astar<commit_after>use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::collections::{BinaryHeap, HashMap};

use std::hash::Hash;

use super::visit::{EdgeRef, GraphBase, IntoEdges, Visitable};
use crate::scored::MinScored;

use crate::algo::Measure;

/// \[Generic\] A* shortest path algorithm.
///
/// Computes the shortest path from `start` to `finish`, including the total path cost.
///
/// `finish` is implicitly given via the `is_goal` callback, which should return `true` if the
/// given node is the finish node.
///
/// The function `edge_cost` should return the cost for a particular edge. Edge costs must be
/// non-negative.
///
/// The function `estimate_cost` should return the estimated cost to the finish for a particular
/// node. For the algorithm to find the actual shortest path, it should be admissible, meaning that
/// it should never overestimate the actual cost to get to the nearest goal node. Estimate costs
/// must also be non-negative.
///
/// The graph should be `Visitable` and implement `IntoEdges`.
///
/// # Example
/// ```
/// use petgraph::Graph;
/// use petgraph::algo::astar;
///
/// let mut g = Graph::new();
/// let a = g.add_node((0., 0.));
/// let b = g.add_node((2., 0.));
/// let c = g.add_node((1., 1.));
/// let d = g.add_node((0., 2.));
/// let e = g.add_node((3., 3.));
/// let f = g.add_node((4., 2.));
/// g.extend_with_edges(&[
///     (a, b, 2),
///     (a, d, 4),
///     (b, c, 1),
///     (b, f, 7),
///     (c, e, 5),
///     (e, f, 1),
///     (d, e, 1),
/// ]);
///
/// // Graph represented with the weight of each edge
/// // Edges with '*' are part of the optimal path.
/// //
/// //     2       1
/// // a ----- b ----- c
/// // | 4*    | 7     |
/// // d       f       | 5
/// // | 1*    | 1*    |
/// // \------ e ------/
///
/// let path = astar(&g, a, |finish| finish == f, |e| *e.weight(), |_| 0);
/// assert_eq!(path, Some((6, vec![a, d, e, f])));
/// ```
///
/// Returns the total cost + the path of subsequent `NodeId` from start to finish, if one was
/// found.
pub fn astar<G, F, H, K, IsGoal>(
    graph: G,
    start: G::NodeId,
    mut is_goal: IsGoal,
    mut edge_cost: F,
    mut estimate_cost: H,
) -> Option<(K, Vec<G::NodeId>)>
where
    G: IntoEdges + Visitable,
    IsGoal: FnMut(G::NodeId) -> bool,
    G::NodeId: Eq + Hash,
    F: FnMut(G::EdgeRef) -> K,
    H: FnMut(G::NodeId) -> K,
    K: Measure + Copy,
{
    let mut visit_next = BinaryHeap::new();
    let mut scores = HashMap::new(); // g-values, cost to reach the node
    let mut estimate_scores = HashMap::new(); // f-values, cost to reach + estimate cost to goal
    let mut path_tracker = PathTracker::<G>::new();

    let zero_score = K::default();
    scores.insert(start, zero_score);
    visit_next.push(MinScored(estimate_cost(start), start));

    while let Some(MinScored(estimate_score, node)) = visit_next.pop() {
        if is_goal(node) {
            let path = path_tracker.reconstruct_path_to(node);
            let cost = scores[&node];
            return Some((cost, path));
        }

        // This lookup can be unwrapped without fear of panic since the node was necessarily scored
        // before adding it to `visit_next`.
        let node_score = scores[&node];

        match estimate_scores.entry(node) {
            Occupied(mut entry) => {
                // If the node has already been visited with an equal or lower score than now, then
                // we do not need to re-visit it.
                if *entry.get() <= estimate_score {
                    continue;
                }
                entry.insert(estimate_score);
            }
            Vacant(entry) => {
                entry.insert(estimate_score);
            }
        }

        for edge in graph.edges(node) {
            let next = edge.target();
            let next_score = node_score + edge_cost(edge);

            match scores.entry(next) {
                Occupied(mut entry) => {
                    // No need to add neighbors that we have already reached through a shorter path
                    // than now.
                    if *entry.get() <= next_score {
                        continue;
                    }
                    entry.insert(next_score);
                }
                Vacant(entry) => {
                    entry.insert(next_score);
                }
            }

            path_tracker.set_predecessor(next, node);
            let next_estimate_score = next_score + estimate_cost(next);
            visit_next.push(MinScored(next_estimate_score, next));
        }
    }

    None
}

struct PathTracker<G>
where
    G: GraphBase,
    G::NodeId: Eq + Hash,
{
    came_from: HashMap<G::NodeId, G::NodeId>,
}

impl<G> PathTracker<G>
where
    G: GraphBase,
    G::NodeId: Eq + Hash,
{
    fn new() -> PathTracker<G> {
        PathTracker {
            came_from: HashMap::new(),
        }
    }

    fn set_predecessor(&mut self, node: G::NodeId, previous: G::NodeId) {
        self.came_from.insert(node, previous);
    }

    fn reconstruct_path_to(&self, last: G::NodeId) -> Vec<G::NodeId> {
        let mut path = vec![last];

        let mut current = last;
        while let Some(&previous) = self.came_from.get(&current) {
            path.push(previous);
            current = previous;
        }

        path.reverse();

        path
    }
}
<|endoftext|>


<commit_before>// Copyright 2015-2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// Inspired by Paul Woolcock's cargo-fmt (https://github.com/pwoolcoc/cargo-fmt/)

#![cfg(not(test))]
#![deny(warnings)]

extern crate getopts;
extern crate serde_json as json;

use std::env;
use std::io::Write;
use std::path::PathBuf;
use std::process::{Command, ExitStatus};
use std::str;
use std::collections::HashSet;
use std::iter::FromIterator;

use getopts::{Matches, Options};
use json::Value;

fn main() {
    let exit_status = execute();
    std::io::stdout().flush().unwrap();
    std::process::exit(exit_status);
}

fn execute() -> i32 {
    let success = 0;
    let failure = 1;

    let mut opts = getopts::Options::new();
    opts.optflag("h", "help", "show this message");
    opts.optflag("q", "quiet", "no output printed to stdout");
    opts.optflag("v", "verbose", "use verbose output");
    opts.optmulti(
        "p",
        "package",
        "specify package to format (only usable in workspaces)",
        "<package>",
    );
    opts.optflag("", "all", "format all packages (only usable in workspaces)");

    // If there is any invalid argument passed to `cargo fmt`, return without formatting.
    if let Some(arg) = env::args()
        .skip(2)
        .take_while(|a| a != "--")
        .find(|a| !a.starts_with('-'))
    {
        print_usage_to_stderr(&opts, &format!("Invalid argument: `{}`.", arg));
        return failure;
    }

    let matches = match opts.parse(env::args().skip(1).take_while(|a| a != "--")) {
        Ok(m) => m,
        Err(e) => {
            print_usage_to_stderr(&opts, &e.to_string());
            return failure;
        }
    };

    let verbosity = match (matches.opt_present("v"), matches.opt_present("q")) {
        (false, false) => Verbosity::Normal,
        (false, true) => Verbosity::Quiet,
        (true, false) => Verbosity::Verbose,
        (true, true) => {
            print_usage_to_stderr(&opts, "quiet mode and verbose mode are not compatible");
            return failure;
        }
    };

    if matches.opt_present("h") {
        print_usage_to_stdout(&opts, "");
        return success;
    }

    let workspace_hitlist = WorkspaceHitlist::from_matches(&matches);

    match format_crate(verbosity, &workspace_hitlist) {
        Err(e) => {
            print_usage_to_stderr(&opts, &e.to_string());
            failure
        }
        Ok(status) => if status.success() {
            success
        } else {
            status.code().unwrap_or(failure)
        },
    }
}

macro_rules! print_usage {
    ($print:ident, $opts:ident, $reason:expr) => ({
        let msg = format!("{}\nusage: cargo fmt [options]", $reason);
        $print!(
            "{}\nThis utility formats all bin and lib files of the current crate using rustfmt. \
             Arguments after `--` are passed to rustfmt.",
            $opts.usage(&msg)
        );
    })
}

fn print_usage_to_stdout(opts: &Options, reason: &str) {
    print_usage!(println, opts, reason);
}

fn print_usage_to_stderr(opts: &Options, reason: &str) {
    print_usage!(eprintln, opts, reason);
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Verbosity {
    Verbose,
    Normal,
    Quiet,
}

fn format_crate(
    verbosity: Verbosity,
    workspace_hitlist: &WorkspaceHitlist,
) -> Result<ExitStatus, std::io::Error> {
    let targets = get_targets(workspace_hitlist)?;

    // Currently only bin and lib files get formatted
    let files: Vec<_> = targets
        .into_iter()
        .filter(|t| t.kind.should_format())
        .inspect(|t| if verbosity == Verbosity::Verbose {
            println!("[{:?}] {:?}", t.kind, t.path)
        })
        .map(|t| t.path)
        .collect();

    format_files(&files, &get_fmt_args(), verbosity)
}

fn get_fmt_args() -> Vec<String> {
    // All arguments after -- are passed to rustfmt
    env::args().skip_while(|a| a != "--").skip(1).collect()
}

#[derive(Debug)]
enum TargetKind {
    Lib,         // dylib, staticlib, lib
    Bin,         // bin
    Example,     // example file
    Test,        // test file
    Bench,       // bench file
    CustomBuild, // build script
    ProcMacro,   // a proc macro implementation
    Other,       // plugin,...
}

impl TargetKind {
    fn should_format(&self) -> bool {
        match *self {
            TargetKind::Lib |
            TargetKind::Bin |
            TargetKind::Example |
            TargetKind::Test |
            TargetKind::Bench |
            TargetKind::CustomBuild |
            TargetKind::ProcMacro => true,
            _ => false,
        }
    }
}

#[derive(Debug)]
pub struct Target {
    path: PathBuf,
    kind: TargetKind,
}

#[derive(Debug, PartialEq, Eq)]
pub enum WorkspaceHitlist {
    All,
    Some(Vec<String>),
    None,
}

impl WorkspaceHitlist {
    pub fn get_some(&self) -> Option<&[String]> {
        if let WorkspaceHitlist::Some(ref hitlist) = *self {
            Some(hitlist)
        } else {
            None
        }
    }

    pub fn from_matches(matches: &Matches) -> WorkspaceHitlist {
        match (matches.opt_present("all"), matches.opt_present("p")) {
            (false, false) => WorkspaceHitlist::None,
            (true, _) => WorkspaceHitlist::All,
            (false, true) => WorkspaceHitlist::Some(matches.opt_strs("p")),
        }
    }
}

// Returns a vector of all compile targets of a crate
fn get_targets(workspace_hitlist: &WorkspaceHitlist) -> Result<Vec<Target>, std::io::Error> {
    let mut targets: Vec<Target> = vec![];

        }
    }
    let output = Command::new("cargo")
        .arg("metadata")
        .arg("--no-deps")
        .output()?;
    if output.status.success() {
        let data = &String::from_utf8(output.stdout).unwrap();
        let json: Value = json::from_str(data).unwrap();
        let json_obj = json.as_object().unwrap();
        let mut hitlist: HashSet<&String> = if *workspace_hitlist != WorkspaceHitlist::All {
            HashSet::from_iter(workspace_hitlist.get_some().unwrap())
        } else {
            HashSet::new() // Unused
        };
        let members: Vec<&Value> = json_obj
            .get("packages")
            .unwrap()
            .as_array()
            .unwrap()
            .into_iter()
            .filter(|member| if *workspace_hitlist == WorkspaceHitlist::All {
                true
            } else {
                let member_obj = member.as_object().unwrap();
                let member_name = member_obj.get("name").unwrap().as_str().unwrap();
                hitlist.take(&member_name.to_string()).is_some()
            })
            .collect();
        if !hitlist.is_empty() {
            // Mimick cargo of only outputting one <package> spec.
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                format!(
                    "package `{}` is not a member of the workspace",
                    hitlist.iter().next().unwrap()
                ),
            ));
        }
        for member in members {
            let member_obj = member.as_object().unwrap();
            let jtargets = member_obj.get("targets").unwrap().as_array().unwrap();
            for jtarget in jtargets {
                targets.push(target_from_json(jtarget));
            }
        }
}

fn target_from_json(jtarget: &Value) -> Target {
    let jtarget = jtarget.as_object().unwrap();
    let path = PathBuf::from(jtarget.get("src_path").unwrap().as_str().unwrap());
    let kinds = jtarget.get("kind").unwrap().as_array().unwrap();
    let kind = match kinds[0].as_str().unwrap() {
        "bin" => TargetKind::Bin,
        "lib" | "dylib" | "staticlib" | "cdylib" | "rlib" => TargetKind::Lib,
        "test" => TargetKind::Test,
        "example" => TargetKind::Example,
        "bench" => TargetKind::Bench,
        "custom-build" => TargetKind::CustomBuild,
        "proc-macro" => TargetKind::ProcMacro,
        _ => TargetKind::Other,
    };

    Target {
        path: path,
        kind: kind,
    }
}

fn format_files(
    files: &[PathBuf],
    fmt_args: &[String],
    verbosity: Verbosity,
) -> Result<ExitStatus, std::io::Error> {
    let stdout = if verbosity == Verbosity::Quiet {
        std::process::Stdio::null()
    } else {
        std::process::Stdio::inherit()
    };
    if verbosity == Verbosity::Verbose {
        print!("rustfmt");
        for a in fmt_args.iter() {
            print!(" {}", a);
        }
        for f in files.iter() {
            print!(" {}", f.display());
        }
        println!("");
    }
    let mut command = Command::new("rustfmt")
        .stdout(stdout)
        .args(files)
        .args(fmt_args)
        .spawn()
        .map_err(|e| match e.kind() {
            std::io::ErrorKind::NotFound => std::io::Error::new(
                std::io::ErrorKind::Other,
                "Could not run rustfmt, please make sure it is in your PATH.",
            ),
            _ => e,
        })?;
    command.wait()
}
<commit_msg>Read local dependencies from cargo metadata<commit_after>// Copyright 2015-2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// Inspired by Paul Woolcock's cargo-fmt (https://github.com/pwoolcoc/cargo-fmt/)

#![cfg(not(test))]
#![deny(warnings)]

extern crate getopts;
extern crate serde_json as json;

use std::env;
use std::io::Write;
use std::path::PathBuf;
use std::process::{Command, ExitStatus};
use std::str;
use std::collections::HashSet;
use std::iter::FromIterator;

use getopts::{Matches, Options};
use json::Value;

fn main() {
    let exit_status = execute();
    std::io::stdout().flush().unwrap();
    std::process::exit(exit_status);
}

fn execute() -> i32 {
    let success = 0;
    let failure = 1;

    let mut opts = getopts::Options::new();
    opts.optflag("h", "help", "show this message");
    opts.optflag("q", "quiet", "no output printed to stdout");
    opts.optflag("v", "verbose", "use verbose output");
    opts.optmulti(
        "p",
        "package",
        "specify package to format (only usable in workspaces)",
        "<package>",
    );
    opts.optflag("", "all", "format all packages (only usable in workspaces)");

    // If there is any invalid argument passed to `cargo fmt`, return without formatting.
    if let Some(arg) = env::args()
        .skip(2)
        .take_while(|a| a != "--")
        .find(|a| !a.starts_with('-'))
    {
        print_usage_to_stderr(&opts, &format!("Invalid argument: `{}`.", arg));
        return failure;
    }

    let matches = match opts.parse(env::args().skip(1).take_while(|a| a != "--")) {
        Ok(m) => m,
        Err(e) => {
            print_usage_to_stderr(&opts, &e.to_string());
            return failure;
        }
    };

    let verbosity = match (matches.opt_present("v"), matches.opt_present("q")) {
        (false, false) => Verbosity::Normal,
        (false, true) => Verbosity::Quiet,
        (true, false) => Verbosity::Verbose,
        (true, true) => {
            print_usage_to_stderr(&opts, "quiet mode and verbose mode are not compatible");
            return failure;
        }
    };

    if matches.opt_present("h") {
        print_usage_to_stdout(&opts, "");
        return success;
    }

    let workspace_hitlist = WorkspaceHitlist::from_matches(&matches);

    match format_crate(verbosity, &workspace_hitlist) {
        Err(e) => {
            print_usage_to_stderr(&opts, &e.to_string());
            failure
        }
        Ok(status) => if status.success() {
            success
        } else {
            status.code().unwrap_or(failure)
        },
    }
}

macro_rules! print_usage {
    ($print:ident, $opts:ident, $reason:expr) => ({
        let msg = format!("{}\nusage: cargo fmt [options]", $reason);
        $print!(
            "{}\nThis utility formats all bin and lib files of the current crate using rustfmt. \
             Arguments after `--` are passed to rustfmt.",
            $opts.usage(&msg)
        );
    })
}

fn print_usage_to_stdout(opts: &Options, reason: &str) {
    print_usage!(println, opts, reason);
}

fn print_usage_to_stderr(opts: &Options, reason: &str) {
    print_usage!(eprintln, opts, reason);
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Verbosity {
    Verbose,
    Normal,
    Quiet,
}

fn format_crate(
    verbosity: Verbosity,
    workspace_hitlist: &WorkspaceHitlist,
) -> Result<ExitStatus, std::io::Error> {
    let targets = get_targets(workspace_hitlist)?;

    // Currently only bin and lib files get formatted
    let files: Vec<_> = targets
        .into_iter()
        .filter(|t| t.kind.should_format())
        .inspect(|t| if verbosity == Verbosity::Verbose {
            println!("[{:?}] {:?}", t.kind, t.path)
        })
        .map(|t| t.path)
        .collect();

    format_files(&files, &get_fmt_args(), verbosity)
}

fn get_fmt_args() -> Vec<String> {
    // All arguments after -- are passed to rustfmt
    env::args().skip_while(|a| a != "--").skip(1).collect()
}

#[derive(Debug)]
enum TargetKind {
    Lib,         // dylib, staticlib, lib
    Bin,         // bin
    Example,     // example file
    Test,        // test file
    Bench,       // bench file
    CustomBuild, // build script
    ProcMacro,   // a proc macro implementation
    Other,       // plugin,...
}

impl TargetKind {
    fn should_format(&self) -> bool {
        match *self {
            TargetKind::Lib |
            TargetKind::Bin |
            TargetKind::Example |
            TargetKind::Test |
            TargetKind::Bench |
            TargetKind::CustomBuild |
            TargetKind::ProcMacro => true,
            _ => false,
        }
    }
}

#[derive(Debug)]
pub struct Target {
    path: PathBuf,
    kind: TargetKind,
}

impl Target {
    pub fn from_json(json_val: &Value) -> Option<Self> {
        let jtarget = json_val.as_object()?;
        let path = PathBuf::from(jtarget.get("src_path")?.as_str()?);
        let kinds = jtarget.get("kind")?.as_array()?;
        let kind = match kinds[0].as_str()? {
            "bin" => TargetKind::Bin,
            "lib" | "dylib" | "staticlib" | "cdylib" | "rlib" => TargetKind::Lib,
            "test" => TargetKind::Test,
            "example" => TargetKind::Example,
            "bench" => TargetKind::Bench,
            "custom-build" => TargetKind::CustomBuild,
            "proc-macro" => TargetKind::ProcMacro,
            _ => TargetKind::Other,
        };

        Some(Target {
            path: path,
            kind: kind,
        })
    }
}

#[derive(Debug, PartialEq, Eq)]
pub enum WorkspaceHitlist {
    All,
    Some(Vec<String>),
    None,
}

impl WorkspaceHitlist {
    pub fn get_some(&self) -> Option<&[String]> {
        if let WorkspaceHitlist::Some(ref hitlist) = *self {
            Some(hitlist)
        } else {
            None
        }
    }

    pub fn from_matches(matches: &Matches) -> WorkspaceHitlist {
        match (matches.opt_present("all"), matches.opt_present("p")) {
            (false, false) => WorkspaceHitlist::None,
            (true, _) => WorkspaceHitlist::All,
            (false, true) => WorkspaceHitlist::Some(matches.opt_strs("p")),
        }
    }
}

fn get_cargo_metadata_from_utf8(v: &[u8]) -> Option<Value> {
    json::from_str(str::from_utf8(v).ok()?).ok()
}

fn get_json_array_with<'a>(v: &'a Value, key: &str) -> Option<&'a Vec<Value>> {
    v.as_object()?.get(key)?.as_array()
}

// `cargo metadata --no-deps | jq '.["packages"]'`
fn get_packages(v: &[u8]) -> Result<Vec<Value>, io::Error> {
    let e = io::Error::new(
        io::ErrorKind::NotFound,
        String::from("`cargo metadata` returned json without a 'packages' key"),
    );
    match get_cargo_metadata_from_utf8(v) {
        Some(ref json_obj) => get_json_array_with(json_obj, "packages").cloned().ok_or(e),
        None => Err(e),
    }
}

fn extract_target_from_package(package: &Value) -> Option<Vec<Target>> {
    let jtargets = get_json_array_with(package, "targets")?;
    let mut targets: Vec<Target> = vec![];
    for jtarget in jtargets {
        targets.push(Target::from_json(&jtarget)?);
    }
    Some(targets)
}

fn filter_packages_with_hitlist<'a>(
    packages: Vec<Value>,
    workspace_hitlist: &'a WorkspaceHitlist,
) -> Result<Vec<Value>, &'a String> {
    if *workspace_hitlist == WorkspaceHitlist::All {
        return Ok(packages);
    }
    let mut hitlist: HashSet<&String> = workspace_hitlist
        .get_some()
        .map_or(HashSet::new(), HashSet::from_iter);
    let members: Vec<Value> = packages
        .into_iter()
        .filter(|member| {
            member
                .as_object()
                .and_then(|member_obj| {
                    member_obj
                        .get("name")
                        .and_then(Value::as_str)
                        .map(|member_name| {
                            hitlist.take(&member_name.to_string()).is_some()
                        })
                })
                .unwrap_or(false)
        })
        .collect();
    if hitlist.is_empty() {
        Ok(members)
    } else {
        Err(hitlist.into_iter().next().unwrap())
    }
}

fn get_dependencies_from_package(package: &Value) -> Option<Vec<PathBuf>> {
    let jdependencies = get_json_array_with(package, "dependencies")?;
    let root_path = env::current_dir().ok()?;
    let mut dependencies: Vec<PathBuf> = vec![];
    for jdep in jdependencies {
        let jdependency = jdep.as_object()?;
        if !jdependency.get("source")?.is_null() {
            continue;
        }
        let name = jdependency.get("name")?.as_str()?;
        let mut path = root_path.clone();
        path.push(&name);
        dependencies.push(path);
    }
    Some(dependencies)
}

// Returns a vector of local dependencies under this crate
fn get_path_to_local_dependencies(packages: &[Value]) -> Vec<PathBuf> {
    let mut local_dependencies: Vec<PathBuf> = vec![];
    for package in packages {
        if let Some(mut d) = get_dependencies_from_package(package) {
            local_dependencies.append(&mut d);
        }
    }
    local_dependencies
}

// Returns a vector of all compile targets of a crate
fn get_targets(workspace_hitlist: &WorkspaceHitlist) -> Result<Vec<Target>, io::Error> {
    let output = Command::new("cargo")
        .args(&["metadata", "--no-deps", "--format-version=1"])
        .output()?;
    if output.status.success() {
        let cur_dir = env::current_dir()?;
        let mut targets: Vec<Target> = vec![];
        let packages = get_packages(&output.stdout)?;

        // If we can find any local dependencies, we will try to get targets from those as well.
        for path in get_path_to_local_dependencies(&packages) {
            env::set_current_dir(path)?;
            targets.append(&mut get_targets(workspace_hitlist)?);
        }

        env::set_current_dir(cur_dir)?;
        match filter_packages_with_hitlist(packages, workspace_hitlist) {
            Ok(packages) => {
                for package in packages {
                    if let Some(mut target) = extract_target_from_package(&package) {
                        targets.append(&mut target);
                    }
                }
                Ok(targets)
            }
            Err(package) => {
                // Mimick cargo of only outputting one <package> spec.
                Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    format!("package `{}` is not a member of the workspace", package),
                ))
            }
        }
    } else {
        Err(io::Error::new(
            io::ErrorKind::NotFound,
            str::from_utf8(&output.stderr).unwrap(),
        ))
    }
}

fn format_files(
    files: &[PathBuf],
    fmt_args: &[String],
    verbosity: Verbosity,
) -> Result<ExitStatus, std::io::Error> {
    let stdout = if verbosity == Verbosity::Quiet {
        std::process::Stdio::null()
    } else {
        std::process::Stdio::inherit()
    };
    if verbosity == Verbosity::Verbose {
        print!("rustfmt");
        for a in fmt_args.iter() {
            print!(" {}", a);
        }
        for f in files.iter() {
            print!(" {}", f.display());
        }
        println!("");
    }
    let mut command = Command::new("rustfmt")
        .stdout(stdout)
        .args(files)
        .args(fmt_args)
        .spawn()
        .map_err(|e| match e.kind() {
            std::io::ErrorKind::NotFound => std::io::Error::new(
                std::io::ErrorKind::Other,
                "Could not run rustfmt, please make sure it is in your PATH.",
            ),
            _ => e,
        })?;
    command.wait()
}
<|endoftext|>


<commit_before>//! Check CPU usage

extern crate rustc_serialize;

extern crate docopt;

use docopt::Docopt;

use std::fs::File;
use std::io::{BufReader,Read};
use std::path::Path;

static USAGE: &'static str = "
Usage: check-cpu [options] [--type=<work-source>]

Options:
    -h, --help             Show this help message

    -s, --sleep=<seconds>  Seconds to collect for [default: 1]
    -w, --warn=<percent>   Percent to warn at     [default: 80]
    -c, --crit=<percent>   Percent to critical at [default: 95]

CPU Work Types:
    Specifying one of the CPU kinds checks that kind of utilization. The
    default is to check total utilization.

    --type=<usage>  One of:
                       total user nice system idle
                       iowait irq softirq steal guest [default: total]
";

#[derive(RustcDecodable, Debug)]
enum WorkSource {
    Total, User, Nice, System, Idle, IoWait, Irq, SoftIrq, Steal, Guest, GuestNice
}

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_help: bool,
    flag_sleep: i32,
    flag_warn:  isize,
    flag_crit:  isize,

    flag_type: WorkSource,
}

/// The number of calculations that have occured on this computer since it
/// started
#[derive(Debug)]
pub struct Calculations {
    pub user: f64,
    pub nice: f64,
    pub system: f64,
    pub idle: f64,
    pub iowait: f64,
    pub irq: f64,
    pub softirq: f64,
    pub steal: f64,
    pub guest: f64,
    pub guest_nice: Option<f64>,
}

impl Calculations {
    /// Create a new `Calculations` struct with the difference between self and other
    pub fn sub(&self, other: &Calculations) -> Calculations {
        Calculations {
            user: self.user - other.user,
            nice: self.nice - other.nice,
            system: self.system - other.system,
            idle: self.idle - other.idle,
            iowait: self.iowait - other.iowait,
            irq: self.irq - other.irq,
            softirq: self.softirq - other.softirq,
            steal: self.steal - other.steal,
            guest: self.guest - other.guest,
            guest_nice: match (self.guest_nice, other.guest_nice) {
                (Some(snice), Some(onice)) => Some(snice - onice),
                _ => None
            }
        }
    }

    /// Jiffies spent non-idle
    ///
    /// This includes all processes in user space, kernel space, and time
    /// stolen by other VMs.
    pub fn active(&self) -> f64 {
        self.user + self.nice // user processes
            + self.system + self.irq + self.softirq // kernel and interrupts
            + self.steal  // vm interrupts
    }

    /// Jiffies spent with nothing to do
    pub fn idle(&self) -> f64 {
        self.idle + self.iowait
    }

    /// Jiffies spent running child VMs
    ///
    /// This is included in `active()`, so don't add this to that when
    /// totalling.
    pub fn virt(&self) -> f64 {
        self.guest + self.guest_nice.unwrap_or(0.0)
    }

    /// All jiffies since the kernel started tracking
    pub fn total(&self) -> f64 {
        self.active() + self.idle()
    }
}

/// Return how much the specific worksource took
///
/// The number returned is between 0 and 100
fn percent_util(kind: &WorkSource, start: &Calculations, end: &Calculations) -> f64 {
    let (start_val, end_val) = match *kind {
        WorkSource::Total => (start.active(), end.active()),
        WorkSource::User => (start.user, end.user),
        WorkSource::IoWait => (start.iowait, end.iowait),
        WorkSource::Nice => (start.nice, end.nice),
        WorkSource::System => (start.system, end.system),
        WorkSource::Idle => (start.idle, end.idle),
        WorkSource::Irq => (start.irq, end.irq),
        WorkSource::SoftIrq => (start.softirq, end.softirq),
        WorkSource::Steal => (start.steal, end.steal),
        WorkSource::Guest => (start.guest, end.guest),
        WorkSource::GuestNice => (start.guest_nice.unwrap_or(0f64), end.guest_nice.unwrap_or(0f64)),
    };
    let total = (end_val - start_val) /
        (end.total() - start.total());
    total * 100f64
}

pub fn read_cpu() -> Calculations {
    let contents = match File::open(&Path::new("/proc/stat")) {
        Ok(ref mut content) => {
            let mut s = String::new();
            let _ = BufReader::new(content).read_to_string(&mut s);
            s
        },
        Err(e) => panic!("Unable to read /proc/stat: {:?}", e)
    };
    let mut word = String::new();
    let mut usages = Vec::new();
    for chr in contents.chars() {
        match chr {
            ' ' => {
                if word != "" && word != "cpu" {
                    let usage = match word.parse() {
                        Ok(num) => num,
                        Err(e) => panic!("Unable to parse '{}' as f64: {:?}", word, e)
                    };
                    usages.push(usage)
                };
                word.clear();
            },
            '\n' => break,
            _ => word.push(chr)
        }
    }

    Calculations {
        user: usages[0],
        nice: usages[1],
        system: usages[2],
        idle: usages[3],
        iowait: usages[4],
        irq: usages[5],
        softirq: usages[6],
        steal: usages[7],
        guest: usages[8],
        guest_nice: match usages.get(9) {
            Some(n) => Some(n.clone()),
            None => None
        },
    }
}

#[cfg_attr(test, allow(dead_code))]
fn main() {
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());
    if args.flag_help {
        print!("{}", USAGE);
        return;
    }

    let start = read_cpu();
    std::thread::sleep_ms((args.flag_sleep * 1000) as u32);
    let end = read_cpu();
    let total = percent_util(&args.flag_type, &start, &end);
    let mut exit_status = 0;
    if total > args.flag_crit as f64 {
        exit_status = 2;
        println!("check-cpu critical: {:?} > {}", total, args.flag_crit);
    } else if total > args.flag_warn as f64 {
        exit_status = 1;
        println!("check-cpu warning: {} > {}", total, args.flag_warn);
    } else {
        println!("check-cpu ok");
    }
    println!("{:?}={:?} ({:?})", args.flag_type, total, end.sub(&start));
    std::process::exit(exit_status);
}

#[cfg(test)]
mod test {
    use docopt::Docopt;
    use super::{USAGE, WorkSource, Calculations, percent_util};

    #[test]
    fn validate_docstring() {
        Docopt::new(USAGE).unwrap();
    }

    fn begintime() -> Calculations {
        Calculations {
            user: 100.0,
            nice: 100.0,
            system: 100.0,
            idle: 100.0,
            iowait: 100.0,
            irq: 100.0,
            softirq: 100.0,
            steal: 100.0,
            guest: 100.0,
            guest_nice: Some(0.0),
        }
    }

    #[test]
    fn percentage_user_idle() {
        let start = begintime();

        let end = Calculations {
            user: 110.0,
            idle: 110.0,
            ..start
        };

        assert_eq!(percent_util(&WorkSource::User, &start, &end), 50.0)
    }

    #[test]
    fn percentage_user() {
        let start = begintime();

        let end = Calculations {
            user: 110.0,
            ..start
        };

        assert_eq!(percent_util(&WorkSource::User, &start, &end), 100.0)
    }

    #[test]
    fn percentage_total_user_idle() {
        let start = begintime();
        let end = Calculations {
            user: 110.0,
            idle: 110.0,
            ..start
        };

        assert_eq!(percent_util(&WorkSource::Total, &start, &end), 50.0)
    }

    #[test]
    fn percentage_total_user_idle_system_steal() {
        let start = begintime();
        let end = Calculations {
            user: 110.0,
            system: 110.0,
            steal: 110.0,
            idle: 110.0,
            ..start
        };

        assert_eq!(percent_util(&WorkSource::Total, &start, &end), 75.0)
    }
}
<commit_msg>c-p: more tests, allow specifying multiple types<commit_after>//! Check CPU usage

extern crate rustc_serialize;

extern crate docopt;

extern crate turbine_plugins;

use docopt::Docopt;
use turbine_plugins::ExitStatus;

use std::fs::File;
use std::io::{BufReader,Read};
use std::path::Path;

static USAGE: &'static str = "
Usage: check-cpu [options] [--type=<work-source>...]

Options:
    -h, --help               Show this help message

    -s, --sample=<seconds>   Seconds to spent collecting   [default: 1]
    -w, --warn=<percent>     Percent to warn at            [default: 80]
    -c, --crit=<percent>     Percent to critical at        [default: 95]

CPU Work Types:

    Specifying one of the CPU kinds checks that kind of utilization. The
    default is to check total utilization.

    --type=<usage>           Some of:
                                total user nice system idle
                                iowait irq softirq steal guest [default: total]
";

#[derive(RustcDecodable, Debug)]
enum WorkSource {
    Total, User, Nice, System, Idle, IoWait, Irq, SoftIrq, Steal, Guest, GuestNice
}

impl std::fmt::Display for WorkSource {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        use WorkSource::*;
        let s = match *self {
            Total => "total",
            User => "user",
            Nice => "nice",
            System => "system",
            Idle => "idle",
            IoWait => "iowait",
            Irq => "irq",
            SoftIrq => "softirq",
            Steal => "steal",
            Guest => "guest",
            GuestNice => "guestnice"
        };
        f.write_str(s)
    }
}

#[derive(RustcDecodable, Debug)]
struct Args {
    flag_help: bool,
    flag_sample: i32,
    flag_warn: isize,
    flag_crit: isize,

    flag_type: Vec<WorkSource>,
}

/// The number of calculations that have occured on this computer since it
/// started
#[derive(Debug)]
struct Calculations {
    user: f64,
    nice: f64,
    system: f64,
    idle: f64,
    iowait: f64,
    irq: f64,
    softirq: f64,
    steal: f64,
    guest: f64,
    guest_nice: Option<f64>,
}

impl Calculations {
    /// Jiffies spent non-idle
    ///
    /// This includes all processes in user space, kernel space, and time
    /// stolen by other VMs.
    fn active(&self) -> f64 {
        self.user + self.nice // user processes
            + self.system + self.irq + self.softirq // kernel and interrupts
            + self.steal  // other VMs stealing our precious time
    }

    /// Jiffies spent with nothing to do
    fn idle(&self) -> f64 {
        self.idle + self.iowait
    }

    /// Jiffies spent running child VMs
    ///
    /// This is included in `active()`, so don't add this to that when
    /// totalling.
    #[allow(dead_code)]  // this mostly exists as documentation of what `guest` means
    fn virt(&self) -> f64 {
        self.guest + self.guest_nice.unwrap_or(0.0)
    }

    /// All jiffies since the kernel started tracking
    fn total(&self) -> f64 {
        self.active() + self.idle()
    }
}

/// Return how much cpu time the specific worksource took
///
/// The number returned is between 0 and 100
fn percent_util(kind: &WorkSource, start: &Calculations, end: &Calculations) -> f64 {
    let (start_val, end_val) = match *kind {
        WorkSource::Total => (start.active(), end.active()),
        WorkSource::User => (start.user, end.user),
        WorkSource::IoWait => (start.iowait, end.iowait),
        WorkSource::Nice => (start.nice, end.nice),
        WorkSource::System => (start.system, end.system),
        WorkSource::Idle => (start.idle, end.idle),
        WorkSource::Irq => (start.irq, end.irq),
        WorkSource::SoftIrq => (start.softirq, end.softirq),
        WorkSource::Steal => (start.steal, end.steal),
        WorkSource::Guest => (start.guest, end.guest),
        WorkSource::GuestNice => (start.guest_nice.unwrap_or(0f64), end.guest_nice.unwrap_or(0f64)),
    };
    let total = (end_val - start_val) /
        (end.total() - start.total());
    total * 100f64
}

fn read_cpu() -> Calculations {
    let contents = match File::open(&Path::new("/proc/stat")) {
        Ok(ref mut content) => {
            let mut s = String::new();
            let _ = BufReader::new(content).read_to_string(&mut s);
            s
        },
        Err(e) => panic!("Unable to read /proc/stat: {:?}", e)
    };
    let mut word = String::new();
    let mut usages = Vec::new();
    for chr in contents.chars() {
        match chr {
            ' ' => {
                if word != "" && word != "cpu" {
                    let usage = match word.parse() {
                        Ok(num) => num,
                        Err(e) => panic!("Unable to parse '{}' as f64: {:?}", word, e)
                    };
                    usages.push(usage)
                };
                word.clear();
            },
            '\n' => break,
            _ => word.push(chr)
        }
    }

    Calculations {
        user: usages[0],
        nice: usages[1],
        system: usages[2],
        idle: usages[3],
        iowait: usages[4],
        irq: usages[5],
        softirq: usages[6],
        steal: usages[7],
        guest: usages[8],
        guest_nice: match usages.get(9) {
            Some(n) => Some(n.clone()),
            None => None
        },
    }
}

fn do_comparison(args: &Args, start: &Calculations, end: &Calculations) -> ExitStatus {
    let mut exit_status = ExitStatus::Ok;

    for flag in &args.flag_type {
        let total = percent_util(flag, &start, &end);
        if total > args.flag_crit as f64 {
            exit_status = std::cmp::max(exit_status, ExitStatus::Critical);
            println!("CRITICAL [check-cpu]: {} {} > {}", flag, total, args.flag_crit);
        } else if total > args.flag_warn as f64 {
            exit_status = std::cmp::max(exit_status, ExitStatus::Warning);
            println!("WARNING [check-cpu]: {} {} > {}", flag, total, args.flag_warn);
        } else {
            println!("OK [check-cpu]");
        }
    }

    exit_status
}

#[cfg_attr(test, allow(dead_code))]
fn main() {
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());
    if args.flag_help {
        print!("{}", USAGE);
        return;
    }

    let start = read_cpu();
    std::thread::sleep_ms((args.flag_sample * 1000) as u32);
    let end = read_cpu();
    let status = do_comparison(&args, &start, &end);
    status.exit()
}

#[cfg(test)]
mod unit {
    use docopt::Docopt;
    use super::{USAGE, WorkSource, Calculations, percent_util};

    #[test]
    fn validate_docstring() {
        Docopt::new(USAGE).unwrap();
    }

    #[test]
    fn validate_allows_multiple_worksources() {
        let argv = || vec!["check-cpu", "--type", "total", "--type", "steal"];
        let _: super::Args = Docopt::new(USAGE)
            .and_then(|d| d.argv(argv().into_iter()).decode())
            .unwrap();
    }

    fn begintime() -> Calculations {
        Calculations {
            user: 100.0,
            nice: 100.0,
            system: 100.0,
            idle: 100.0,
            iowait: 100.0,
            irq: 100.0,
            softirq: 100.0,
            steal: 100.0,
            guest: 100.0,
            guest_nice: Some(0.0),
        }
    }

    #[test]
    fn percentage_user_idle() {
        let start = begintime();

        let end = Calculations {
            user: 110.0,
            idle: 110.0,
            ..start
        };

        assert_eq!(percent_util(&WorkSource::User, &start, &end), 50.0)
    }

    #[test]
    fn percentage_user() {
        let start = begintime();

        let end = Calculations {
            user: 110.0,
            ..start
        };

        assert_eq!(percent_util(&WorkSource::User, &start, &end), 100.0)
    }

    #[test]
    fn percentage_total_user_idle() {
        let start = begintime();
        let end = Calculations {
            user: 110.0,
            idle: 110.0,
            ..start
        };

        assert_eq!(percent_util(&WorkSource::Total, &start, &end), 50.0)
    }

    #[test]
    fn percentage_total_user_idle_system_steal() {
        let start = begintime();
        let end = Calculations {
            user: 110.0,
            system: 110.0,
            steal: 110.0,
            idle: 110.0,
            ..start
        };

        assert_eq!(percent_util(&WorkSource::Total, &start, &end), 75.0)
    }
}

#[cfg(test)]
mod integration {
    // not really integration tests, but higher level
    use super::{do_comparison, Calculations, USAGE};

    use turbine_plugins::ExitStatus;
    use docopt::Docopt;

    fn start() -> Calculations {
        Calculations {
            user: 100.0,
            nice: 100.0,
            system: 100.0,
            idle: 100.0,
            iowait: 100.0,
            irq: 100.0,
            softirq: 100.0,
            steal: 100.0,
            guest: 100.0,
            guest_nice: Some(0.0),
        }
    }

    #[test]
    fn does_alert() {
        let argv = || vec!["check-cpu", "-c", "49", "--type", "total", "--type", "steal"];
        let start = start();
        let end = Calculations {
            user: 110.0,
            idle: 110.0,
            ..start
        };
        let args: super::Args = Docopt::new(USAGE)
            .and_then(|d| d.argv(argv().into_iter()).decode())
            .unwrap();

        assert_eq!(do_comparison(&args, &start, &end),
                   ExitStatus::Critical);
    }
}
<|endoftext|>


<commit_before>//! This module provides usefull traits that where deprecated in rust

// Note copied from the stdlib under MIT license

use std::ops::{ Add, Div, Mul, Rem, Sub };
use num::{ Float, NumCast };
use std::marker::Reflect;

/// Num trait from old stdlib
pub trait Num: Zero + One + Add<Output=Self> + Sub<Output=Self> + Mul<Output=Self> + Div<Output=Self> + Rem<Output=Self> + PartialEq<Self> { }
impl<A: Zero + One + Add<Output=A> + Sub<Output=A> + Mul<Output=A> + Div<Output=A> + Rem<Output=A> + PartialEq<A>> Num for A { }

/// Zero trait from old stdlib
pub trait Zero: Add<Output=Self> {
    /// Returns the zero value for T
    fn zero() -> Self;
    /// Returns true if zero.
    fn is_zero(&self) -> bool;
}

/// Returns the zero value for T
#[allow(unused)]
pub fn zero<T: Zero>() -> T { Zero::zero() }
macro_rules! zero_impl {
    ($t:ty, $v:expr) => {
        impl Zero for $t {
            fn zero() -> $t { $v }
            fn is_zero(&self) -> bool { *self == $v }
        }
    }
}
zero_impl!(usize, 0usize);
zero_impl!(u8,   0u8);
zero_impl!(u16,  0u16);
zero_impl!(u32,  0u32);
zero_impl!(u64,  0u64);
zero_impl!(isize, 0isize);
zero_impl!(i8,  0i8);
zero_impl!(i16, 0i16);
zero_impl!(i32, 0i32);
zero_impl!(i64, 0i64);
zero_impl!(f32, 0.0f32);
zero_impl!(f64, 0.0f64);

/// Primitive trait from old stdlib, added max_value
pub trait Primitive: Copy + NumCast + Num + PartialOrd<Self> + Clone + Reflect {
    /// The maximum value of primitive.
    fn max_value() -> Self;
}

impl Primitive for usize {
    fn max_value() -> usize { usize::max_value() }
}
impl Primitive for u8 {
    fn max_value() -> u8 { u8::max_value() }
}
impl Primitive for u16 {
    fn max_value() -> u16 { u16::max_value() }
}
impl Primitive for u32 {
    fn max_value() -> u32 { u32::max_value() }
}
impl Primitive for u64 {
    fn max_value() -> u64 { u64::max_value() }
}
impl Primitive for isize {
    fn max_value() -> isize { isize::max_value() }
}
impl Primitive for i8 {
    fn max_value() -> i8 { i8::max_value() }
}
impl Primitive for i16 {
    fn max_value() -> i16 { i16::max_value() }
}
impl Primitive for i32 {
    fn max_value() -> i32 { i32::max_value() }
}
impl Primitive for i64 {
    fn max_value() -> i64 { i64::max_value() }
}
impl Primitive for f32 {
    fn max_value() -> f32 { Float::max_value() }
}
impl Primitive for f64 {
    fn max_value() -> f64 { Float::max_value() }
}

/// One trait from old stdlib, added max_value
pub trait One: Mul<Output=Self> {
    /// Returns the unit value of Self
    fn one() -> Self;
}
/// Returns the unit value of T
#[allow(unused)]
pub fn one<T: One>() -> T { One::one() }
macro_rules! one_impl {
    ($t:ty, $v:expr) => {
        impl One for $t {
            fn one() -> $t { $v }
        }
    }
}
one_impl!(usize, 1usize);
one_impl!(u8,  1u8);
one_impl!(u16, 1u16);
one_impl!(u32, 1u32);
one_impl!(u64, 1u64);
one_impl!(isize, 1isize);
one_impl!(i8,  1i8);
one_impl!(i16, 1i16);
one_impl!(i32, 1i32);
one_impl!(i64, 1i64);
one_impl!(f32, 1.0f32);
one_impl!(f64, 1.0f64);
<commit_msg>Replace traits::One with num::One<commit_after>//! This module provides usefull traits that where deprecated in rust

// Note copied from the stdlib under MIT license

use std::ops::{ Add, Div, Mul, Rem, Sub };
use num::{ Float, NumCast, One };
use std::marker::Reflect;

/// Num trait from old stdlib
pub trait Num: Zero + One + Add<Output=Self> + Sub<Output=Self> + Mul<Output=Self> + Div<Output=Self> + Rem<Output=Self> + PartialEq<Self> { }
impl<A: Zero + One + Add<Output=A> + Sub<Output=A> + Mul<Output=A> + Div<Output=A> + Rem<Output=A> + PartialEq<A>> Num for A { }

/// Zero trait from old stdlib
pub trait Zero: Add<Output=Self> {
    /// Returns the zero value for T
    fn zero() -> Self;
    /// Returns true if zero.
    fn is_zero(&self) -> bool;
}

/// Returns the zero value for T
#[allow(unused)]
pub fn zero<T: Zero>() -> T { Zero::zero() }
macro_rules! zero_impl {
    ($t:ty, $v:expr) => {
        impl Zero for $t {
            fn zero() -> $t { $v }
            fn is_zero(&self) -> bool { *self == $v }
        }
    }
}
zero_impl!(usize, 0usize);
zero_impl!(u8,   0u8);
zero_impl!(u16,  0u16);
zero_impl!(u32,  0u32);
zero_impl!(u64,  0u64);
zero_impl!(isize, 0isize);
zero_impl!(i8,  0i8);
zero_impl!(i16, 0i16);
zero_impl!(i32, 0i32);
zero_impl!(i64, 0i64);
zero_impl!(f32, 0.0f32);
zero_impl!(f64, 0.0f64);

/// Primitive trait from old stdlib, added max_value
pub trait Primitive: Copy + NumCast + Num + PartialOrd<Self> + Clone + Reflect {
    /// The maximum value of primitive.
    fn max_value() -> Self;
}

impl Primitive for usize {
    fn max_value() -> usize { usize::max_value() }
}
impl Primitive for u8 {
    fn max_value() -> u8 { u8::max_value() }
}
impl Primitive for u16 {
    fn max_value() -> u16 { u16::max_value() }
}
impl Primitive for u32 {
    fn max_value() -> u32 { u32::max_value() }
}
impl Primitive for u64 {
    fn max_value() -> u64 { u64::max_value() }
}
impl Primitive for isize {
    fn max_value() -> isize { isize::max_value() }
}
impl Primitive for i8 {
    fn max_value() -> i8 { i8::max_value() }
}
impl Primitive for i16 {
    fn max_value() -> i16 { i16::max_value() }
}
impl Primitive for i32 {
    fn max_value() -> i32 { i32::max_value() }
}
impl Primitive for i64 {
    fn max_value() -> i64 { i64::max_value() }
}
impl Primitive for f32 {
    fn max_value() -> f32 { Float::max_value() }
}
impl Primitive for f64 {
    fn max_value() -> f64 { Float::max_value() }
}
<|endoftext|>


<commit_before>use super::{completer::IonCompleter, InteractiveBinary};
use std::io::ErrorKind;

pub fn readln(binary: &InteractiveBinary) -> Option<String> {
    let shell = binary.shell.borrow();
    let mut completer = IonCompleter::new(&shell);
    let prompt = binary.prompt();
    let line = binary.context.borrow_mut().read_line(prompt, None, &mut completer);

    match line {
        Ok(line) => {
            if line.bytes().next() != Some(b'#') && line.bytes().any(|c| !c.is_ascii_whitespace()) {
                binary.shell.borrow_mut().unterminated = true;
            }
            Some(line)
        }
        // Handles Ctrl + C
        Err(ref err) if err.kind() == ErrorKind::Interrupted => None,
        // Handles Ctrl + D
        Err(ref err) if err.kind() == ErrorKind::UnexpectedEof => {
            let mut shell = binary.shell.borrow_mut();
            if !shell.unterminated && shell.exit_block().is_err() {
                shell.exit(None);
            }
            None
        }
        Err(err) => {
            eprintln!("ion: liner: {}", err);
            None
        }
    }
}
<commit_msg>Fix the panic on start<commit_after>use super::{completer::IonCompleter, InteractiveBinary};
use std::io::ErrorKind;

pub fn readln(binary: &InteractiveBinary) -> Option<String> {
    let prompt = binary.prompt();
    let line = binary.context.borrow_mut().read_line(
        prompt,
        None,
        &mut IonCompleter::new(&binary.shell.borrow()),
    );

    match line {
        Ok(line) => {
            if line.bytes().next() != Some(b'#') && line.bytes().any(|c| !c.is_ascii_whitespace()) {
                binary.shell.borrow_mut().unterminated = true;
            }
            Some(line)
        }
        // Handles Ctrl + C
        Err(ref err) if err.kind() == ErrorKind::Interrupted => None,
        // Handles Ctrl + D
        Err(ref err) if err.kind() == ErrorKind::UnexpectedEof => {
            let mut shell = binary.shell.borrow_mut();
            if !shell.unterminated && shell.exit_block().is_err() {
                shell.exit(None);
            }
            None
        }
        Err(err) => {
            eprintln!("ion: liner: {}", err);
            None
        }
    }
}
<|endoftext|>


<commit_before>extern crate luthor;

use std::rc::Rc;
use std::cell::RefCell;
use std::old_io::{File, Open, Read, Write};
use std::old_io::IoError;
use std::old_io::IoResult;
use super::GapBuffer;
use super::gap_buffer;
use super::Position;
use super::Range;
use super::Cursor;
use self::luthor::token::{Token, Category};

/// A UTF-8 buffer with bounds-checked cursor management and persistence.
pub struct Buffer {
    data: Rc<RefCell<GapBuffer>>,
    lexer: Option<fn(&str) -> Vec<Token>>,
    pub path: Option<Path>,
    pub cursor: Cursor,
}

impl Buffer {
    /// Returns the contents of the buffer as a string.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// assert_eq!(buffer.data(), "scribe");
    /// ```
    pub fn data(&self) -> String {
        self.data.borrow().to_string()
    }

    /// Writes the contents of the buffer to its path.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::old_io::{File, Open, Read};
    ///
    /// // Set up a buffer and point it to a path.
    /// let mut buffer = scribe::buffer::new();
    /// let write_path = Path::new("my_doc");
    /// buffer.path = Some(write_path.clone());
    ///
    /// // Put some data into the buffer and save it.
    /// buffer.insert("scribe");
    /// buffer.save();
    ///
    /// # let saved_data = File::open_mode(&write_path, Open, Read)
    /// #   .unwrap().read_to_string().unwrap();
    /// # assert_eq!(saved_data, "scribe");
    ///
    /// # std::old_io::fs::unlink(&write_path);
    /// ```
    pub fn save(&self) -> Option<IoError> {
        let path = match self.path.clone() {
            Some(p) => p,
            None => Path::new(""),
        };

        // Try to open and write to the file, returning any errors encountered.
        let mut file = match File::open_mode(&path, Open, Write) {
            Ok(f) => f,
            Err(error) => return Some(error),
        };

        // We use to_string here because we don't want to write the gap contents.
        match file.write(self.data().to_string().as_bytes()) {
            Ok(_) => (),
            Err(error) => return Some(error),
        }

        return None
    }

    /// Inserts `data` into the buffer at the cursor position.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// assert_eq!(buffer.data(), "scribe");
    /// ```
    pub fn insert(&mut self, data: &str) {
        self.data.borrow_mut().insert(data, &self.cursor);
    }

    /// Deletes a character at the cursor position.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// buffer.delete();
    /// assert_eq!(buffer.data(), "cribe");
    /// ```
    pub fn delete(&mut self) {
        let mut end = self.cursor.clone();
        end.offset += 1;
        self.data.borrow_mut().delete(&Range{ start: *self.cursor, end: end});
    }

    /// Produces a set of tokens based on the buffer data
    /// suitable for colorized display, using a lexer for the
    /// buffer data's language and/or format. If a lexer is not
    /// available, the set will consist of a single text-category token.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    ///
    /// // Build the buffer data string back by combining its token lexemes.
    /// let mut data = String::new();
    /// for token in buffer.tokens().iter() {
    ///     data.push_str(token.lexeme.as_slice());
    /// }
    /// assert_eq!(data, "scribe");
    /// ```
    pub fn tokens(&self) -> Vec<Token> {
        match self.lexer {
            Some(lexer) => lexer(&self.data()),
            None => vec![Token{ lexeme: self.data(), category: Category::Text }],
        }
    }
}

/// Creates a new empty buffer. The buffer's cursor is set to the beginning of the buffer.
///
/// # Examples
///
/// ```
/// let buffer = scribe::buffer::new();
/// # assert_eq!(buffer.cursor.line, 0);
/// # assert_eq!(buffer.cursor.offset, 0);
/// ```
pub fn new() -> Buffer {
    let data = Rc::new(RefCell::new(gap_buffer::new(String::new())));
    let cursor = Cursor{ data: data.clone(), position: Position{ line: 0, offset: 0 }};

    Buffer{ data: data.clone(), path: None, cursor: cursor, lexer: None }
}

/// Creates a new buffer by reading the UTF-8 interpreted file contents of the specified path.
/// The buffer's cursor is set to the beginning of the buffer.
///
/// # Examples
///
/// ```
/// let buffer = scribe::buffer::from_file(&Path::new("tests/sample/file")).unwrap();
/// assert_eq!(buffer.data(), "it works!\n");
/// # assert_eq!(buffer.cursor.line, 0);
/// # assert_eq!(buffer.cursor.offset, 0);
/// ```
pub fn from_file(path: &Path) -> IoResult<Buffer> {
    // Try to open and read the file, returning any errors encountered.
    let mut file = match File::open_mode(path, Open, Read) {
        Ok(f) => f,
        Err(error) => return Err(error),
    };
    let mut data = match file.read_to_string() {
        Ok(d) => d,
        Err(error) => return Err(error),
    };

    let data = Rc::new(RefCell::new(gap_buffer::new(data)));
    let cursor = Cursor{ data: data.clone(), position: Position{ line: 0, offset: 0 }};

    // Create a new buffer using the loaded data, path, and other defaults.
    Ok(Buffer{ data: data.clone(), path: Some(path.clone()), cursor: cursor, lexer: None })
}

#[cfg(test)]
mod tests {
    use super::new;
    use super::luthor::token::{Token, Category};

    fn placeholder_lexer(_: &str) -> Vec<Token> {
        vec![Token{ lexeme: "lexer".to_string(), category: Category::Text }]
    }

    #[test]
    fn tokens_returns_one_text_token_when_no_lexer_is_set() {
        let mut buffer = new();
        buffer.insert("scribe");
        let expected_tokens = vec![Token{ lexeme: "scribe".to_string(), category: Category::Text }];
        assert_eq!(buffer.tokens(), expected_tokens);
    }

    #[test]
    fn tokens_returns_result_of_lexer_when_set() {
        let mut buffer = new();
        let expected_tokens = placeholder_lexer("scribe");
        buffer.lexer = Some(placeholder_lexer);
        assert_eq!(buffer.tokens(), expected_tokens);
    }
}
<commit_msg>Use type detection to implement JSON lexer configuration.<commit_after>extern crate luthor;

use std::rc::Rc;
use std::cell::RefCell;
use std::old_io::{File, Open, Read, Write};
use std::old_io::IoError;
use std::old_io::IoResult;
use super::GapBuffer;
use super::gap_buffer;
use super::Position;
use super::Range;
use super::Cursor;
use super::type_detection;
use self::luthor::token::{Token, Category};
use self::luthor::lexers;

/// A UTF-8 buffer with bounds-checked cursor management and persistence.
pub struct Buffer {
    data: Rc<RefCell<GapBuffer>>,
    lexer: Option<fn(&str) -> Vec<Token>>,
    pub path: Option<Path>,
    pub cursor: Cursor,
}

impl Buffer {
    /// Returns the contents of the buffer as a string.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// assert_eq!(buffer.data(), "scribe");
    /// ```
    pub fn data(&self) -> String {
        self.data.borrow().to_string()
    }

    /// Writes the contents of the buffer to its path.
    ///
    /// # Examples
    ///
    /// ```
    /// # use std::old_io::{File, Open, Read};
    ///
    /// // Set up a buffer and point it to a path.
    /// let mut buffer = scribe::buffer::new();
    /// let write_path = Path::new("my_doc");
    /// buffer.path = Some(write_path.clone());
    ///
    /// // Put some data into the buffer and save it.
    /// buffer.insert("scribe");
    /// buffer.save();
    ///
    /// # let saved_data = File::open_mode(&write_path, Open, Read)
    /// #   .unwrap().read_to_string().unwrap();
    /// # assert_eq!(saved_data, "scribe");
    ///
    /// # std::old_io::fs::unlink(&write_path);
    /// ```
    pub fn save(&self) -> Option<IoError> {
        let path = match self.path.clone() {
            Some(p) => p,
            None => Path::new(""),
        };

        // Try to open and write to the file, returning any errors encountered.
        let mut file = match File::open_mode(&path, Open, Write) {
            Ok(f) => f,
            Err(error) => return Some(error),
        };

        // We use to_string here because we don't want to write the gap contents.
        match file.write(self.data().to_string().as_bytes()) {
            Ok(_) => (),
            Err(error) => return Some(error),
        }

        return None
    }

    /// Inserts `data` into the buffer at the cursor position.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// assert_eq!(buffer.data(), "scribe");
    /// ```
    pub fn insert(&mut self, data: &str) {
        self.data.borrow_mut().insert(data, &self.cursor);
    }

    /// Deletes a character at the cursor position.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    /// buffer.delete();
    /// assert_eq!(buffer.data(), "cribe");
    /// ```
    pub fn delete(&mut self) {
        let mut end = self.cursor.clone();
        end.offset += 1;
        self.data.borrow_mut().delete(&Range{ start: *self.cursor, end: end});
    }

    /// Produces a set of tokens based on the buffer data
    /// suitable for colorized display, using a lexer for the
    /// buffer data's language and/or format. If a lexer is not
    /// available, the set will consist of a single text-category token.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut buffer = scribe::buffer::new();
    /// buffer.insert("scribe");
    ///
    /// // Build the buffer data string back by combining its token lexemes.
    /// let mut data = String::new();
    /// for token in buffer.tokens().iter() {
    ///     data.push_str(token.lexeme.as_slice());
    /// }
    /// assert_eq!(data, "scribe");
    /// ```
    pub fn tokens(&self) -> Vec<Token> {
        match self.lexer {
            Some(lexer) => lexer(&self.data()),
            None => vec![Token{ lexeme: self.data(), category: Category::Text }],
        }
    }
}

/// Creates a new empty buffer. The buffer's cursor is set to the beginning of the buffer.
///
/// # Examples
///
/// ```
/// let buffer = scribe::buffer::new();
/// # assert_eq!(buffer.cursor.line, 0);
/// # assert_eq!(buffer.cursor.offset, 0);
/// ```
pub fn new() -> Buffer {
    let data = Rc::new(RefCell::new(gap_buffer::new(String::new())));
    let cursor = Cursor{ data: data.clone(), position: Position{ line: 0, offset: 0 }};

    Buffer{ data: data.clone(), path: None, cursor: cursor, lexer: None }
}

/// Creates a new buffer by reading the UTF-8 interpreted file contents of the specified path.
/// The buffer's cursor is set to the beginning of the buffer. The buffer data's type will be
/// inferred based on its extension, and an appropriate lexer will be used, if available (see
/// tokens method for further information on why this happens).
///
/// # Examples
///
/// ```
/// let buffer = scribe::buffer::from_file(&Path::new("tests/sample/file")).unwrap();
/// assert_eq!(buffer.data(), "it works!\n");
/// # assert_eq!(buffer.cursor.line, 0);
/// # assert_eq!(buffer.cursor.offset, 0);
/// ```
pub fn from_file(path: &Path) -> IoResult<Buffer> {
    // Try to open and read the file, returning any errors encountered.
    let mut file = match File::open_mode(path, Open, Read) {
        Ok(f) => f,
        Err(error) => return Err(error),
    };
    let mut data = match file.read_to_string() {
        Ok(d) => d,
        Err(error) => return Err(error),
    };

    let data = Rc::new(RefCell::new(gap_buffer::new(data)));
    let cursor = Cursor{ data: data.clone(), position: Position{ line: 0, offset: 0 }};

    // Detect the file type and use its corresponding lexer, if available.
    let lexer = match type_detection::from_path(path) {
        Some(type_detection::Type::JSON) => Some(lexers::json::lex as fn(&str) -> Vec<Token>),
        _ => None,
    };

    // Create a new buffer using the loaded data, path, and other defaults.
    Ok(Buffer{ data: data.clone(), path: Some(path.clone()), cursor: cursor, lexer: lexer })
}

#[cfg(test)]
mod tests {
    use super::new;
    use super::luthor::token::{Token, Category};

    fn placeholder_lexer(_: &str) -> Vec<Token> {
        vec![Token{ lexeme: "lexer".to_string(), category: Category::Text }]
    }

    #[test]
    fn tokens_returns_one_text_token_when_no_lexer_is_set() {
        let mut buffer = new();
        buffer.insert("scribe");
        let expected_tokens = vec![Token{ lexeme: "scribe".to_string(), category: Category::Text }];
        assert_eq!(buffer.tokens(), expected_tokens);
    }

    #[test]
    fn tokens_returns_result_of_lexer_when_set() {
        let mut buffer = new();
        let expected_tokens = placeholder_lexer("scribe");
        buffer.lexer = Some(placeholder_lexer);
        assert_eq!(buffer.tokens(), expected_tokens);
    }
}
<|endoftext|>