<commit_before>use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::io;
use std::rc::Rc;
use std::{i32, u64};

use bytes::BytesMut;
use futures::sync::mpsc;
use futures::{self, Async, Poll, StartSend, Stream};
use tokio_core::reactor::Handle;
use tokio_io::codec::{Decoder, Encoder, Framed};
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_proto::multiplex::{ClientProto, RequestId};

use asnom::common::TagClass;
use asnom::{Consumer, ConsumerState, Input, IResult, Move};
use asnom::structure::{PL, StructureTag};
use asnom::structures::{ASNTag, Integer, Null, Sequence, Tag};
use asnom::parse::parse_uint;
use asnom::parse::Parser;
use asnom::universal::Types;
use asnom::write;

use ldap::LdapOp;
use search::SearchItem;

pub type LdapRequestId = i32;

pub struct ProtoBundle {
    pub search_helpers: HashMap<RequestId, SearchHelper>,
    pub id_map: HashMap<LdapRequestId, RequestId>,
    pub next_id: LdapRequestId,
    pub handle: Handle,
}

impl ProtoBundle {
    fn create_search_helper(&mut self, id: RequestId, tx: mpsc::UnboundedSender<SearchItem>) {
        self.search_helpers.insert(id, SearchHelper {
            seen: false,
            msgid: 0,           // not valid, must be properly initialized later
            tx: tx,
        });
    }

    fn inc_next_id(&mut self) -> LdapRequestId {
        if self.next_id == i32::MAX {
            self.next_id = 0;
        }
        self.next_id += 1;
        self.next_id
    }
}

pub struct SearchHelper {
    pub seen: bool,
    pub msgid: LdapRequestId,
    pub tx: mpsc::UnboundedSender<SearchItem>,
}

impl SearchHelper {
    fn send_item(&mut self, item: SearchItem) -> io::Result<()> {
        self.tx.send(item).map_err(|e| io::Error::new(io::ErrorKind::Other, format!("{:?}", e)))
    }
}

#[derive(Debug)]
pub struct LdapResult {
    pub rc: u8,
    pub matched: String,
    pub text: String,
    pub refs: Vec<HashSet<String>>,
}

impl From<Tag> for LdapResult {
    fn from(t: Tag) -> LdapResult {
        match t {
            Tag::StructureTag(t) => t.into(),
            _ => unimplemented!(),
        }
    }
}

impl From<StructureTag> for LdapResult {
    fn from(t: StructureTag) -> LdapResult {
        let mut tags = t.expect_constructed().expect("result sequence").into_iter();
        let rc = match parse_uint(tags.next().expect("element")
                .match_class(TagClass::Universal)
                .and_then(|t| t.match_id(Types::Enumerated as u64))
                .and_then(|t| t.expect_primitive()).expect("result code")
                .as_slice()) {
            IResult::Done(_, rc) => rc as u8,
            _ => panic!("failed to parse result code"),
        };
        let matched = String::from_utf8(tags.next().expect("element").expect_primitive().expect("octet string"))
            .expect("matched dn");
        let text = String::from_utf8(tags.next().expect("element").expect_primitive().expect("octet string"))
            .expect("diagnostic message");
        let mut refs = Vec::new();
        match tags.next() {
            None => (),
            Some(raw_refs) => {
                let raw_refs = match raw_refs.match_class(TagClass::Context)
                        .and_then(|t| t.match_id(3))
                        .and_then(|t| t.expect_constructed()) {
                    Some(rr) => rr,
                    None => panic!("failed to parse referrals"),
                };
                refs.push(raw_refs.into_iter()
                    .map(|t| t.expect_primitive().expect("octet string"))
                    .map(String::from_utf8)
                    .map(|s| s.expect("uri"))
                    .collect());
            },
        }
        LdapResult {
            rc: rc,
            matched: matched,
            text: text,
            refs: refs,
        }
    }
}

pub struct LdapCodec {
    bundle: Rc<RefCell<ProtoBundle>>,
}

impl Decoder for LdapCodec {
    type Item = (RequestId, (Tag, Option<StructureTag>));
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        let decoding_error = io::Error::new(io::ErrorKind::Other, "Decoding error");
        let mut parser = Parser::new();
        let (amt, tag) = match *parser.handle(Input::Element(buf)) {
            ConsumerState::Continue(_) => return Ok(None),
            ConsumerState::Error(_e) => return Err(decoding_error),
            ConsumerState::Done(amt, ref tag) => (amt, tag),
        };
        let amt = match amt {
            Move::Await(_) => return Ok(None),
            Move::Seek(_) => return Err(decoding_error),
            Move::Consume(amt) => amt,
        };
        buf.split_to(amt);
        let tag = tag.clone();
        let mut tags = match tag.match_id(Types::Sequence as u64).and_then(|t| t.expect_constructed()) {
            Some(tags) => tags,
            None => return Err(decoding_error),
        };
        let maybe_controls = tags.pop().expect("element");
        let has_controls = match maybe_controls {
            StructureTag { id, class, ref payload } if class == TagClass::Context && id == 0 => match *payload {
                PL::C(_) => true,
                PL::P(_) => return Err(decoding_error),
            },
            _ => false,
        };
        let (protoop, controls) = if has_controls {
            (tags.pop().expect("element"), Some(maybe_controls))
        } else {
            (maybe_controls, None)
        };
        let msgid = match parse_uint(tags.pop().expect("element")
                .match_class(TagClass::Universal)
                .and_then(|t| t.match_id(Types::Integer as u64))
                .and_then(|t| t.expect_primitive()).expect("message id")
                .as_slice()) {
            IResult::Done(_, id) => id as i32,
            _ => return Err(decoding_error),
        };
        let id = match self.bundle.borrow().id_map.get(&msgid) {
            Some(&id) => id,
            None => return Err(io::Error::new(io::ErrorKind::Other, format!("No id found for message id: {}", msgid))),
        };
        match protoop.id {
            op_id @ 4 | op_id @ 5 | op_id @ 19 => {
                let null_tag = Tag::Null(Null { ..Default::default() });
                let id_tag = Tag::Integer(Integer {
                    inner: id as i64,
                    .. Default::default()
                });
                let mut bundle = self.bundle.borrow_mut();
                let mut helper = match bundle.search_helpers.get_mut(&id) {
                    Some(h) => h,
                    None => return Err(io::Error::new(io::ErrorKind::Other, format!("Id mismatch: {}", id))),
                };
                helper.send_item(match op_id {
                    4 => SearchItem::Entry(protoop),
                    5 => SearchItem::Done(id, protoop.into(), controls),
                    19 => SearchItem::Referral(protoop),
                    _ => panic!("impossible op_id"),
                })?;
                if helper.seen {
                    Ok(Some((u64::MAX, (null_tag, None))))
                } else {
                    helper.seen = true;
                    Ok(Some((id, (id_tag, None))))
                }
            },
            _ => Ok(Some((id, (Tag::StructureTag(protoop), controls)))),
        }
    }
}

impl Encoder for LdapCodec {
    type Item = (RequestId, LdapOp);
    type Error = io::Error;

    fn encode(&mut self, msg: Self::Item, into: &mut BytesMut) -> io::Result<()> {
        let (id, op) = msg;
        let tag = match op {
            LdapOp::Single(tag) => tag,
            LdapOp::Multi(tag, tx) => {
                self.bundle.borrow_mut().create_search_helper(id, tx);
                tag
            },
            _ => unimplemented!(),
        };
        let outtag = {
            // tokio-proto ids are u64, and LDAP (client) message ids are i32 > 0,
            // so we must have wraparound logic and a mapping from the latter to
            // the former
            let mut bundle = self.bundle.borrow_mut();
            let prev_ldap_id = bundle.next_id;
            let mut next_ldap_id = prev_ldap_id;
            while bundle.id_map.entry(next_ldap_id).or_insert(id) != &id {
                next_ldap_id = bundle.inc_next_id();
                assert_ne!(next_ldap_id, prev_ldap_id, "LDAP message id wraparound with no free slots");
            }
            bundle.inc_next_id();
            match bundle.search_helpers.get_mut(&id) {
                Some(ref mut helper) => helper.msgid = next_ldap_id,
                None => (),
            }
            Tag::Sequence(Sequence {
                inner: vec![
                    Tag::Integer(Integer {
                        inner: next_ldap_id as i64,
                        .. Default::default()
                    }),
                    tag,
                ],
                .. Default::default()
            })
        };
        let outstruct = outtag.into_structure();
        trace!("Sending packet: {:?}", &outstruct);
        write::encode_into(into, outstruct)?;
        Ok(())
    }
}

pub struct LdapProto {
    bundle: Rc<RefCell<ProtoBundle>>,
}

impl LdapProto {
    pub fn new(handle: Handle) -> LdapProto {
        LdapProto {
            bundle: Rc::new(RefCell::new(ProtoBundle {
                search_helpers: HashMap::new(),
                id_map: HashMap::new(),
                next_id: 1,
                handle: handle,
            }))
        }
    }

    pub fn bundle(&self) -> Rc<RefCell<ProtoBundle>> {
        self.bundle.clone()
    }
}

pub struct ResponseFilter<T> {
    upstream: T,
}

impl<T> Stream for ResponseFilter<T>
    where T: Stream<Item=(RequestId, (Tag, Option<StructureTag>)), Error=io::Error>
{
    type Item = (RequestId, (Tag, Option<StructureTag>));
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        loop {
            match try_ready!(self.upstream.poll()) {
                Some((id, _)) if id == u64::MAX => continue,
                msg => return Ok(Async::Ready(msg)),
            }
        }
    }
}

impl<T> futures::Sink for ResponseFilter<T>
    where T: futures::Sink<SinkItem=(RequestId, LdapOp), SinkError=io::Error>
{
    type SinkItem = (RequestId, LdapOp);
    type SinkError = io::Error;

    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        self.upstream.poll_complete()
    }
}

impl<T: AsyncRead + AsyncWrite + 'static> ClientProto<T> for LdapProto {
    type Request = LdapOp;
    type Response = (Tag, Option<StructureTag>);

    type Transport = ResponseFilter<Framed<T, LdapCodec>>;
    type BindTransport = Result<Self::Transport, io::Error>;

    fn bind_transport(&self, io: T) -> Self::BindTransport {
        let ldapcodec = LdapCodec {
            bundle: self.bundle.clone(),
        };
        Ok(ResponseFilter {
            upstream: io.framed(ldapcodec),
        })
    }
}
<commit_msg>Consistently use lower case in error messages<commit_after>use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::io;
use std::rc::Rc;
use std::{i32, u64};

use bytes::BytesMut;
use futures::sync::mpsc;
use futures::{self, Async, Poll, StartSend, Stream};
use tokio_core::reactor::Handle;
use tokio_io::codec::{Decoder, Encoder, Framed};
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_proto::multiplex::{ClientProto, RequestId};

use asnom::common::TagClass;
use asnom::{Consumer, ConsumerState, Input, IResult, Move};
use asnom::structure::{PL, StructureTag};
use asnom::structures::{ASNTag, Integer, Null, Sequence, Tag};
use asnom::parse::parse_uint;
use asnom::parse::Parser;
use asnom::universal::Types;
use asnom::write;

use ldap::LdapOp;
use search::SearchItem;

pub type LdapRequestId = i32;

pub struct ProtoBundle {
    pub search_helpers: HashMap<RequestId, SearchHelper>,
    pub id_map: HashMap<LdapRequestId, RequestId>,
    pub next_id: LdapRequestId,
    pub handle: Handle,
}

impl ProtoBundle {
    fn create_search_helper(&mut self, id: RequestId, tx: mpsc::UnboundedSender<SearchItem>) {
        self.search_helpers.insert(id, SearchHelper {
            seen: false,
            msgid: 0,           // not valid, must be properly initialized later
            tx: tx,
        });
    }

    fn inc_next_id(&mut self) -> LdapRequestId {
        if self.next_id == i32::MAX {
            self.next_id = 0;
        }
        self.next_id += 1;
        self.next_id
    }
}

pub struct SearchHelper {
    pub seen: bool,
    pub msgid: LdapRequestId,
    pub tx: mpsc::UnboundedSender<SearchItem>,
}

impl SearchHelper {
    fn send_item(&mut self, item: SearchItem) -> io::Result<()> {
        self.tx.send(item).map_err(|e| io::Error::new(io::ErrorKind::Other, format!("{:?}", e)))
    }
}

#[derive(Debug)]
pub struct LdapResult {
    pub rc: u8,
    pub matched: String,
    pub text: String,
    pub refs: Vec<HashSet<String>>,
}

impl From<Tag> for LdapResult {
    fn from(t: Tag) -> LdapResult {
        match t {
            Tag::StructureTag(t) => t.into(),
            _ => unimplemented!(),
        }
    }
}

impl From<StructureTag> for LdapResult {
    fn from(t: StructureTag) -> LdapResult {
        let mut tags = t.expect_constructed().expect("result sequence").into_iter();
        let rc = match parse_uint(tags.next().expect("element")
                .match_class(TagClass::Universal)
                .and_then(|t| t.match_id(Types::Enumerated as u64))
                .and_then(|t| t.expect_primitive()).expect("result code")
                .as_slice()) {
            IResult::Done(_, rc) => rc as u8,
            _ => panic!("failed to parse result code"),
        };
        let matched = String::from_utf8(tags.next().expect("element").expect_primitive().expect("octet string"))
            .expect("matched dn");
        let text = String::from_utf8(tags.next().expect("element").expect_primitive().expect("octet string"))
            .expect("diagnostic message");
        let mut refs = Vec::new();
        match tags.next() {
            None => (),
            Some(raw_refs) => {
                let raw_refs = match raw_refs.match_class(TagClass::Context)
                        .and_then(|t| t.match_id(3))
                        .and_then(|t| t.expect_constructed()) {
                    Some(rr) => rr,
                    None => panic!("failed to parse referrals"),
                };
                refs.push(raw_refs.into_iter()
                    .map(|t| t.expect_primitive().expect("octet string"))
                    .map(String::from_utf8)
                    .map(|s| s.expect("uri"))
                    .collect());
            },
        }
        LdapResult {
            rc: rc,
            matched: matched,
            text: text,
            refs: refs,
        }
    }
}

pub struct LdapCodec {
    bundle: Rc<RefCell<ProtoBundle>>,
}

impl Decoder for LdapCodec {
    type Item = (RequestId, (Tag, Option<StructureTag>));
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        let decoding_error = io::Error::new(io::ErrorKind::Other, "decoding error");
        let mut parser = Parser::new();
        let (amt, tag) = match *parser.handle(Input::Element(buf)) {
            ConsumerState::Continue(_) => return Ok(None),
            ConsumerState::Error(_e) => return Err(decoding_error),
            ConsumerState::Done(amt, ref tag) => (amt, tag),
        };
        let amt = match amt {
            Move::Await(_) => return Ok(None),
            Move::Seek(_) => return Err(decoding_error),
            Move::Consume(amt) => amt,
        };
        buf.split_to(amt);
        let tag = tag.clone();
        let mut tags = match tag.match_id(Types::Sequence as u64).and_then(|t| t.expect_constructed()) {
            Some(tags) => tags,
            None => return Err(decoding_error),
        };
        let maybe_controls = tags.pop().expect("element");
        let has_controls = match maybe_controls {
            StructureTag { id, class, ref payload } if class == TagClass::Context && id == 0 => match *payload {
                PL::C(_) => true,
                PL::P(_) => return Err(decoding_error),
            },
            _ => false,
        };
        let (protoop, controls) = if has_controls {
            (tags.pop().expect("element"), Some(maybe_controls))
        } else {
            (maybe_controls, None)
        };
        let msgid = match parse_uint(tags.pop().expect("element")
                .match_class(TagClass::Universal)
                .and_then(|t| t.match_id(Types::Integer as u64))
                .and_then(|t| t.expect_primitive()).expect("message id")
                .as_slice()) {
            IResult::Done(_, id) => id as i32,
            _ => return Err(decoding_error),
        };
        let id = match self.bundle.borrow().id_map.get(&msgid) {
            Some(&id) => id,
            None => return Err(io::Error::new(io::ErrorKind::Other, format!("no id found for message id: {}", msgid))),
        };
        match protoop.id {
            op_id @ 4 | op_id @ 5 | op_id @ 19 => {
                let null_tag = Tag::Null(Null { ..Default::default() });
                let id_tag = Tag::Integer(Integer {
                    inner: id as i64,
                    .. Default::default()
                });
                let mut bundle = self.bundle.borrow_mut();
                let mut helper = match bundle.search_helpers.get_mut(&id) {
                    Some(h) => h,
                    None => return Err(io::Error::new(io::ErrorKind::Other, format!("id mismatch: {}", id))),
                };
                helper.send_item(match op_id {
                    4 => SearchItem::Entry(protoop),
                    5 => SearchItem::Done(id, protoop.into(), controls),
                    19 => SearchItem::Referral(protoop),
                    _ => panic!("impossible op_id"),
                })?;
                if helper.seen {
                    Ok(Some((u64::MAX, (null_tag, None))))
                } else {
                    helper.seen = true;
                    Ok(Some((id, (id_tag, None))))
                }
            },
            _ => Ok(Some((id, (Tag::StructureTag(protoop), controls)))),
        }
    }
}

impl Encoder for LdapCodec {
    type Item = (RequestId, LdapOp);
    type Error = io::Error;

    fn encode(&mut self, msg: Self::Item, into: &mut BytesMut) -> io::Result<()> {
        let (id, op) = msg;
        let tag = match op {
            LdapOp::Single(tag) => tag,
            LdapOp::Multi(tag, tx) => {
                self.bundle.borrow_mut().create_search_helper(id, tx);
                tag
            },
            _ => unimplemented!(),
        };
        let outtag = {
            // tokio-proto ids are u64, and LDAP (client) message ids are i32 > 0,
            // so we must have wraparound logic and a mapping from the latter to
            // the former
            let mut bundle = self.bundle.borrow_mut();
            let prev_ldap_id = bundle.next_id;
            let mut next_ldap_id = prev_ldap_id;
            while bundle.id_map.entry(next_ldap_id).or_insert(id) != &id {
                next_ldap_id = bundle.inc_next_id();
                assert_ne!(next_ldap_id, prev_ldap_id, "LDAP message id wraparound with no free slots");
            }
            bundle.inc_next_id();
            match bundle.search_helpers.get_mut(&id) {
                Some(ref mut helper) => helper.msgid = next_ldap_id,
                None => (),
            }
            Tag::Sequence(Sequence {
                inner: vec![
                    Tag::Integer(Integer {
                        inner: next_ldap_id as i64,
                        .. Default::default()
                    }),
                    tag,
                ],
                .. Default::default()
            })
        };
        let outstruct = outtag.into_structure();
        trace!("Sending packet: {:?}", &outstruct);
        write::encode_into(into, outstruct)?;
        Ok(())
    }
}

pub struct LdapProto {
    bundle: Rc<RefCell<ProtoBundle>>,
}

impl LdapProto {
    pub fn new(handle: Handle) -> LdapProto {
        LdapProto {
            bundle: Rc::new(RefCell::new(ProtoBundle {
                search_helpers: HashMap::new(),
                id_map: HashMap::new(),
                next_id: 1,
                handle: handle,
            }))
        }
    }

    pub fn bundle(&self) -> Rc<RefCell<ProtoBundle>> {
        self.bundle.clone()
    }
}

pub struct ResponseFilter<T> {
    upstream: T,
}

impl<T> Stream for ResponseFilter<T>
    where T: Stream<Item=(RequestId, (Tag, Option<StructureTag>)), Error=io::Error>
{
    type Item = (RequestId, (Tag, Option<StructureTag>));
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        loop {
            match try_ready!(self.upstream.poll()) {
                Some((id, _)) if id == u64::MAX => continue,
                msg => return Ok(Async::Ready(msg)),
            }
        }
    }
}

impl<T> futures::Sink for ResponseFilter<T>
    where T: futures::Sink<SinkItem=(RequestId, LdapOp), SinkError=io::Error>
{
    type SinkItem = (RequestId, LdapOp);
    type SinkError = io::Error;

    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        self.upstream.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        self.upstream.poll_complete()
    }
}

impl<T: AsyncRead + AsyncWrite + 'static> ClientProto<T> for LdapProto {
    type Request = LdapOp;
    type Response = (Tag, Option<StructureTag>);

    type Transport = ResponseFilter<Framed<T, LdapCodec>>;
    type BindTransport = Result<Self::Transport, io::Error>;

    fn bind_transport(&self, io: T) -> Self::BindTransport {
        let ldapcodec = LdapCodec {
            bundle: self.bundle.clone(),
        };
        Ok(ResponseFilter {
            upstream: io.framed(ldapcodec),
        })
    }
}
<|endoftext|>


<commit_before>extern crate procrs;
extern crate argparse;
#[macro_use]
extern crate prettytable;
use prettytable::Table;
use prettytable::format::FormatBuilder;
use std::collections::HashMap;
use std::iter::repeat;
use procrs::pid::*;
use procrs::TaskId;
use argparse::{ArgumentParser, StoreTrue, Store};

fn main() {
    let opts = parse_args();
    let (query, long, perf, verbose, tree) =
        (opts.query, opts.long, opts.perf, opts.verbose, opts.tree);

    let mut pids = PidIter::new_query(query).unwrap()
        .collect::<Result<Vec<_>, _>>().unwrap();

    let mut name_indent = HashMap::new();

    if verbose {
        for pid in pids {
            println!("{:?}", pid);
        }
        return
    }

    if opts.tree {
        pids = treeify_names(pids, &mut name_indent);
    } else {
        pids.sort_by(|p1, p2| p1.stat.pid.cmp(&p2.stat.pid));
    };

    let mut table = Table::init(
        pids.iter().map(|p| {
            let mut name = match tree {
                false => String::new(),
                true => name_indent.remove(&p.stat.pid).unwrap()
            };

            // TODO: There is a lot of boiler plate for Option<num...>,
            // this can probably be reduced with a custom type (row! calls to_string).
            match (long, perf) {
                (false, false) => {
                    name.push_str(&p.stat.comm);
                    row![p.stat.pid, p.stat.ppid, name]
                },
                (true, false) => {
                    name.push_str(
                        match p.cmdline.join(" ") {
                            ref s if s.len() > 0 => s,
                            _ => &p.stat.comm
                        }
                    );
                    row![p.stat.pid, p.stat.ppid,  name]
                },
                (false, true) => {
                    name.push_str(&p.stat.comm);
                    row![p.stat.pid, p.stat.ppid,
                        p.status.vmrss.map(|m| (m / 1024).to_string()).unwrap_or("".to_owned()),
                        p.stat.utime,
                        name]
                },
                (true, true) => {
                    name.push_str(
                        match p.cmdline.join(" ") {
                            ref s if s.len() > 0 => s,
                            _ => &p.stat.comm
                        }
                    );
                    row![p.stat.pid, p.stat.ppid,
                        p.status.vmrss.map(|m| (m / 1024).to_string()).unwrap_or("".to_owned()),
                        p.stat.utime,
                        name]
                }
            }
        }).collect::<Vec<_>>()
    );

    table.set_titles(
        match (long, perf) {
            (false, false) =>
                row!["Pid", "Ppid", "Cmd"],
            (true, false) =>
                row!["Pid", "Ppid", "Cmd"],
            (false, true) =>
                row!["Pid", "Ppid", "RSS", "Time", "Cmd"],
            (true, true) =>
                row!["Pid", "Ppid", "RSS", "Time", "Cmd"]
        }
    );
    table.set_format(
        FormatBuilder::new()
            .column_separator(' ')
            .build()
    );
    table.printstd();
}

// Given a vector of Pid structs, treeify their names, and return them in the right order.
// This is similar to ps -AH.
fn treeify_names(pids: Vec<Pid>, name_indents: &mut HashMap<TaskId, String>) -> Vec<Pid> {
    let mut child_pids = HashMap::new();
    for pid in pids {
        let ppid = pid.stat.ppid;
        child_pids.entry(ppid)
            .or_insert(Vec::new())
            .push(pid);
    }
    enumerate_children(0, &mut child_pids, name_indents, -1)
}

// Enumerate children pids, and return them.
fn enumerate_children(pid: TaskId, child_pids: &mut HashMap<TaskId, Vec<Pid>>,
    name_indents: &mut HashMap<TaskId, String>, indent: i32) -> Vec<Pid> {
    name_indents.insert(pid,
        match indent {
            i if i >= 0 =>
                repeat("  ").take(i as usize).collect::<String>(),
            _ => "".to_owned()
        }
    );
    let mut pids = Vec::new();
    let ppids = match child_pids.remove(&pid) {
        Some(v) => v,
        None => { return pids; }
    };
    for pid in ppids {
        let pid_num = pid.stat.pid;
        pids.push(pid);
        pids.append(
            &mut enumerate_children(pid_num, child_pids, name_indents, indent + 1)
        );
    }
    pids
}

struct ProgOpts {
    query: PidQuery,
    tree: bool,
    perf: bool,
    long: bool,
    verbose: bool
}

fn parse_args() -> ProgOpts {
    let mut opts = ProgOpts {
        query: PidQuery::NoneQuery,
        tree: false,
        perf: false,
        long: false,
        verbose: false
    };

    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Query linux processes");
        ap.refer(&mut opts.tree)
            .add_option(&["-t", "--tree"], StoreTrue, "Display process tree");
        ap.refer(&mut opts.perf)
            .add_option(&["-p", "--perf"], StoreTrue, "Display performance information");
        ap.refer(&mut opts.long)
            .add_option(&["-l", "--long"], StoreTrue, "Display more information");
        ap.refer(&mut opts.verbose)
            .add_option(&["-v", "--verbose"], StoreTrue, "Verbose output");
        ap.refer(&mut opts.query)
            .add_argument("query", Store, "Optional query to search by, pid or string");
        ap.parse_args_or_exit();
    }

    opts
}
<commit_msg>Improve psq rss and utime output<commit_after>extern crate procrs;
extern crate argparse;
#[macro_use]
extern crate prettytable;
use prettytable::Table;
use prettytable::format::FormatBuilder;
use std::collections::HashMap;
use std::iter::repeat;
use procrs::pid::*;
use procrs::TaskId;
use argparse::{ArgumentParser, StoreTrue, Store};

fn main() {
    let opts = parse_args();
    let (query, long, perf, verbose, tree) =
        (opts.query, opts.long, opts.perf, opts.verbose, opts.tree);

    let mut pids = PidIter::new_query(query).unwrap()
        .collect::<Result<Vec<_>, _>>().unwrap();

    let mut name_indent = HashMap::new();

    if verbose {
        for pid in pids {
            println!("{:?}", pid);
        }
        return
    }

    if opts.tree {
        pids = treeify_names(pids, &mut name_indent);
    } else {
        pids.sort_by(|p1, p2| p1.stat.pid.cmp(&p2.stat.pid));
    };
    // Assume hertz is 100.
    // TODO: Look this up via syscall (no /proc value for it)
    let hertz = 100;
    let minute_hertz = hertz * 60;
    let hour_hertz = minute_hertz * 60;

    let mut table = Table::init(
        pids.iter().map(|p| {
            // When we have a tree, the name is prepended with an indent.
            let mut name = match tree {
                false => String::new(),
                true => name_indent.remove(&p.stat.pid).unwrap()
            };

            // For long output, try using the cmdline first.
            // FIXME: Sometimes prog_name != cmdline[0].
            if !long {
                name.push_str(&p.stat.comm);
            } else {
                let cmdline = p.cmdline.join(" ");
                name.push_str(
                    match cmdline {
                         ref s if s.len() > 0 => s,
                        _ => &p.stat.comm
                    }
                );
            }

            match (long, perf) {
                (_, false) => {
                    row![p.stat.pid, p.stat.ppid, name]
                },
                (_, true) => {
                    let rss = p.status.vmrss.map(|m| (m / 1024).to_string()).unwrap_or("".to_owned());
                    // FIXME: This algorithm is definitely wrong
                    let second_utime = p.stat.utime % hertz;
                    let minute_utime = p.stat.utime / minute_hertz;
                    let hour_utime = p.stat.utime / hour_hertz;
                    let utime = format!(
                        "{:02}:{:02}:{:02}",
                        hour_utime,
                        minute_utime,
                        second_utime
                    );
                    row![p.stat.pid, p.stat.ppid, rss, utime, name]
                }
            }
        }).collect::<Vec<_>>()
    );

    table.set_titles(
        match (long, perf) {
            (false, false) =>
                row!["Pid", "Ppid", "Cmd"],
            (true, false) =>
                row!["Pid", "Ppid", "Cmd"],
            (false, true) =>
                row!["Pid", "Ppid", "RSS", "Time", "Cmd"],
            (true, true) =>
                row!["Pid", "Ppid", "RSS", "Time", "Cmd"]
        }
    );
    table.set_format(
        FormatBuilder::new()
            .column_separator(' ')
            .build()
    );
    table.printstd();
}

// Given a vector of Pid structs, treeify their names, and return them in the right order.
// This is similar to ps -AH.
fn treeify_names(pids: Vec<Pid>, name_indents: &mut HashMap<TaskId, String>) -> Vec<Pid> {
    let mut child_pids = HashMap::new();
    for pid in pids {
        let ppid = pid.stat.ppid;
        child_pids.entry(ppid)
            .or_insert(Vec::new())
            .push(pid);
    }
    enumerate_children(0, &mut child_pids, name_indents, -1)
}

// Enumerate children pids, and return them.
fn enumerate_children(pid: TaskId, child_pids: &mut HashMap<TaskId, Vec<Pid>>,
    name_indents: &mut HashMap<TaskId, String>, indent: i32) -> Vec<Pid> {
    name_indents.insert(pid,
        match indent {
            i if i >= 0 =>
                repeat("  ").take(i as usize).collect::<String>(),
            _ => "".to_owned()
        }
    );
    let mut pids = Vec::new();
    let ppids = match child_pids.remove(&pid) {
        Some(v) => v,
        None => { return pids; }
    };
    for pid in ppids {
        let pid_num = pid.stat.pid;
        pids.push(pid);
        pids.append(
            &mut enumerate_children(pid_num, child_pids, name_indents, indent + 1)
        );
    }
    pids
}

struct ProgOpts {
    query: PidQuery,
    tree: bool,
    perf: bool,
    long: bool,
    verbose: bool
}

fn parse_args() -> ProgOpts {
    let mut opts = ProgOpts {
        query: PidQuery::NoneQuery,
        tree: false,
        perf: false,
        long: false,
        verbose: false
    };

    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Query linux processes");
        ap.refer(&mut opts.tree)
            .add_option(&["-t", "--tree"], StoreTrue, "Display process tree");
        ap.refer(&mut opts.perf)
            .add_option(&["-p", "--perf"], StoreTrue, "Display performance information");
        ap.refer(&mut opts.long)
            .add_option(&["-l", "--long"], StoreTrue, "Display more information");
        ap.refer(&mut opts.verbose)
            .add_option(&["-v", "--verbose"], StoreTrue, "Verbose output");
        ap.refer(&mut opts.query)
            .add_argument("query", Store, "Optional query to search by, pid or string");
        ap.parse_args_or_exit();
    }

    opts
}
<|endoftext|>


<commit_before>//! Advertizes the capabilities of the LSP Server.
use std::env;

use crate::semantic_tokens;

use lsp_types::{
    CallHierarchyServerCapability, CodeActionOptions, CodeActionProviderCapability,
    CodeLensOptions, CompletionOptions, DocumentOnTypeFormattingOptions,
    FoldingRangeProviderCapability, ImplementationProviderCapability, RenameOptions,
    RenameProviderCapability, SaveOptions, SelectionRangeProviderCapability,
    SemanticTokensDocumentProvider, SemanticTokensLegend, SemanticTokensOptions,
    ServerCapabilities, SignatureHelpOptions, TextDocumentSyncCapability, TextDocumentSyncKind,
    TextDocumentSyncOptions, TypeDefinitionProviderCapability, WorkDoneProgressOptions,
};

pub fn server_capabilities() -> ServerCapabilities {
    ServerCapabilities {
        text_document_sync: Some(TextDocumentSyncCapability::Options(TextDocumentSyncOptions {
            open_close: Some(true),
            change: Some(if env::var("RA_PROFILE").is_ok() {
                TextDocumentSyncKind::Incremental
            } else {
                TextDocumentSyncKind::Full
            }),
            will_save: None,
            will_save_wait_until: None,
            save: Some(SaveOptions::default()),
        })),
        hover_provider: Some(true),
        completion_provider: Some(CompletionOptions {
            resolve_provider: None,
            trigger_characters: Some(vec![":".to_string(), ".".to_string()]),
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        }),
        signature_help_provider: Some(SignatureHelpOptions {
            trigger_characters: Some(vec!["(".to_string(), ",".to_string()]),
            retrigger_characters: None,
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        }),
        declaration_provider: None,
        definition_provider: Some(true),
        type_definition_provider: Some(TypeDefinitionProviderCapability::Simple(true)),
        implementation_provider: Some(ImplementationProviderCapability::Simple(true)),
        references_provider: Some(true),
        document_highlight_provider: Some(true),
        document_symbol_provider: Some(true),
        workspace_symbol_provider: Some(true),
        code_action_provider: Some(CodeActionProviderCapability::Options(CodeActionOptions {
            // Advertise support for all built-in CodeActionKinds
            code_action_kinds: Some(vec![
                lsp_types::code_action_kind::EMPTY.to_string(),
                lsp_types::code_action_kind::QUICKFIX.to_string(),
                lsp_types::code_action_kind::REFACTOR.to_string(),
                lsp_types::code_action_kind::REFACTOR_EXTRACT.to_string(),
                lsp_types::code_action_kind::REFACTOR_INLINE.to_string(),
                lsp_types::code_action_kind::REFACTOR_REWRITE.to_string(),
                lsp_types::code_action_kind::SOURCE.to_string(),
                lsp_types::code_action_kind::SOURCE_ORGANIZE_IMPORTS.to_string(),
            ]),
            work_done_progress_options: Default::default(),
        })),
        code_lens_provider: Some(CodeLensOptions { resolve_provider: Some(true) }),
        document_formatting_provider: Some(true),
        document_range_formatting_provider: None,
        document_on_type_formatting_provider: Some(DocumentOnTypeFormattingOptions {
            first_trigger_character: "=".to_string(),
            more_trigger_character: Some(vec![".".to_string(), ">".to_string()]),
        }),
        selection_range_provider: Some(SelectionRangeProviderCapability::Simple(true)),
        semantic_highlighting: None,
        folding_range_provider: Some(FoldingRangeProviderCapability::Simple(true)),
        rename_provider: Some(RenameProviderCapability::Options(RenameOptions {
            prepare_provider: Some(true),
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        })),
        document_link_provider: None,
        color_provider: None,
        execute_command_provider: None,
        workspace: None,
        call_hierarchy_provider: Some(CallHierarchyServerCapability::Simple(true)),
        semantic_tokens_provider: Some(
            SemanticTokensOptions {
                legend: SemanticTokensLegend {
                    token_types: semantic_tokens::SUPPORTED_TYPES.to_vec(),
                    token_modifiers: semantic_tokens::SUPPORTED_MODIFIERS.to_vec(),
                },

                document_provider: Some(SemanticTokensDocumentProvider::Bool(true)),
                range_provider: Some(true),
                work_done_progress_options: Default::default(),
            }
            .into(),
        ),
        experimental: Default::default(),
    }
}
<commit_msg>Make incremental sync opt-out<commit_after>//! Advertizes the capabilities of the LSP Server.
use std::env;

use crate::semantic_tokens;

use lsp_types::{
    CallHierarchyServerCapability, CodeActionOptions, CodeActionProviderCapability,
    CodeLensOptions, CompletionOptions, DocumentOnTypeFormattingOptions,
    FoldingRangeProviderCapability, ImplementationProviderCapability, RenameOptions,
    RenameProviderCapability, SaveOptions, SelectionRangeProviderCapability,
    SemanticTokensDocumentProvider, SemanticTokensLegend, SemanticTokensOptions,
    ServerCapabilities, SignatureHelpOptions, TextDocumentSyncCapability, TextDocumentSyncKind,
    TextDocumentSyncOptions, TypeDefinitionProviderCapability, WorkDoneProgressOptions,
};

pub fn server_capabilities() -> ServerCapabilities {
    ServerCapabilities {
        text_document_sync: Some(TextDocumentSyncCapability::Options(TextDocumentSyncOptions {
            open_close: Some(true),
            change: Some(if env::var("RA_NO_INCREMENTAL_SYNC").is_ok() {
                TextDocumentSyncKind::Full
            } else {
                TextDocumentSyncKind::Incremental
            }),
            will_save: None,
            will_save_wait_until: None,
            save: Some(SaveOptions::default()),
        })),
        hover_provider: Some(true),
        completion_provider: Some(CompletionOptions {
            resolve_provider: None,
            trigger_characters: Some(vec![":".to_string(), ".".to_string()]),
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        }),
        signature_help_provider: Some(SignatureHelpOptions {
            trigger_characters: Some(vec!["(".to_string(), ",".to_string()]),
            retrigger_characters: None,
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        }),
        declaration_provider: None,
        definition_provider: Some(true),
        type_definition_provider: Some(TypeDefinitionProviderCapability::Simple(true)),
        implementation_provider: Some(ImplementationProviderCapability::Simple(true)),
        references_provider: Some(true),
        document_highlight_provider: Some(true),
        document_symbol_provider: Some(true),
        workspace_symbol_provider: Some(true),
        code_action_provider: Some(CodeActionProviderCapability::Options(CodeActionOptions {
            // Advertise support for all built-in CodeActionKinds
            code_action_kinds: Some(vec![
                lsp_types::code_action_kind::EMPTY.to_string(),
                lsp_types::code_action_kind::QUICKFIX.to_string(),
                lsp_types::code_action_kind::REFACTOR.to_string(),
                lsp_types::code_action_kind::REFACTOR_EXTRACT.to_string(),
                lsp_types::code_action_kind::REFACTOR_INLINE.to_string(),
                lsp_types::code_action_kind::REFACTOR_REWRITE.to_string(),
                lsp_types::code_action_kind::SOURCE.to_string(),
                lsp_types::code_action_kind::SOURCE_ORGANIZE_IMPORTS.to_string(),
            ]),
            work_done_progress_options: Default::default(),
        })),
        code_lens_provider: Some(CodeLensOptions { resolve_provider: Some(true) }),
        document_formatting_provider: Some(true),
        document_range_formatting_provider: None,
        document_on_type_formatting_provider: Some(DocumentOnTypeFormattingOptions {
            first_trigger_character: "=".to_string(),
            more_trigger_character: Some(vec![".".to_string(), ">".to_string()]),
        }),
        selection_range_provider: Some(SelectionRangeProviderCapability::Simple(true)),
        semantic_highlighting: None,
        folding_range_provider: Some(FoldingRangeProviderCapability::Simple(true)),
        rename_provider: Some(RenameProviderCapability::Options(RenameOptions {
            prepare_provider: Some(true),
            work_done_progress_options: WorkDoneProgressOptions { work_done_progress: None },
        })),
        document_link_provider: None,
        color_provider: None,
        execute_command_provider: None,
        workspace: None,
        call_hierarchy_provider: Some(CallHierarchyServerCapability::Simple(true)),
        semantic_tokens_provider: Some(
            SemanticTokensOptions {
                legend: SemanticTokensLegend {
                    token_types: semantic_tokens::SUPPORTED_TYPES.to_vec(),
                    token_modifiers: semantic_tokens::SUPPORTED_MODIFIERS.to_vec(),
                },

                document_provider: Some(SemanticTokensDocumentProvider::Bool(true)),
                range_provider: Some(true),
                work_done_progress_options: Default::default(),
            }
            .into(),
        ),
        experimental: Default::default(),
    }
}
<|endoftext|>


<commit_before>use chrono::{Datelike, DateTime, NaiveDate, UTC};
use diesel;
use diesel::expression::dsl::*;
use diesel::prelude::*;
use hyper::client::Client;
use hyper::status::StatusCode;

use DB_POOL;
use domain::releases::Release;
use error::DashResult;

lazy_static! {
    static ref CLIENT: Client = Client::new();
}

pub fn most_recent_update() -> DashResult<DateTime<UTC>> {
    info!("finding most recent nightly release updates");

    let conn = try!(DB_POOL.get());

    let most_recent: NaiveDate = {
        use domain::schema::release::dsl::*;
        try!(release.select(max(date)).filter(released).first(&*conn))
    };

    Ok(DateTime::from_utc(most_recent.and_hms(0, 0, 0), UTC))
}

fn get_release_for_date(d: NaiveDate) -> DashResult<Release> {
    let url = format!("https://static.rust-lang.org/dist/{}-{:02}-{:02}/index.html",
                      d.year(),
                      d.month(),
                      d.day());

    let response = try!(CLIENT.get(&url).send());
    match response.status {
        StatusCode::Ok => Ok(Release { date: d, released: true }),
        _ => Ok(Release { date: d, released: false }),
    }
}

pub fn get_releases_since(d: NaiveDate) -> DashResult<Vec<Release>> {
    let mut releases = vec![];

    let mut curr = d;
    let today = UTC::today().naive_utc();

    while curr <= today {
        let curr_release = try!(get_release_for_date(curr));
        releases.push(curr_release);
        curr = curr.succ();
    }

    Ok(releases)
}

pub fn ingest_releases_since(d: DateTime<UTC>) -> DashResult<()> {
    use diesel::prelude::*;
    use domain::schema::release::dsl::*;
    let releases = try!(get_releases_since(d.date().naive_utc()));

    let conn = try!(DB_POOL.get());

    for r in releases {
        let pk = release.filter(date.eq(r.date))
            .first::<(i32, NaiveDate, bool)>(&*conn)
            .map(|f| f.0)
            .ok();

        if let Some(pk) = pk {
            try!(diesel::update(release.find(pk)).set(&r).execute(&*conn));
        } else {
            try!(diesel::insert(&r).into(release).execute(&*conn));
        }
    }

    Ok(())
}
<commit_msg>Changing the URL used for determining whether a nightly released.<commit_after>use chrono::{Datelike, DateTime, NaiveDate, UTC};
use diesel;
use diesel::expression::dsl::*;
use diesel::prelude::*;
use hyper::client::Client;
use hyper::status::StatusCode;

use DB_POOL;
use domain::releases::Release;
use error::DashResult;

lazy_static! {
    static ref CLIENT: Client = Client::new();
}

pub fn most_recent_update() -> DashResult<DateTime<UTC>> {
    info!("finding most recent nightly release updates");

    let conn = try!(DB_POOL.get());

    let most_recent: NaiveDate = {
        use domain::schema::release::dsl::*;
        try!(release.select(max(date)).filter(released).first(&*conn))
    };

    Ok(DateTime::from_utc(most_recent.and_hms(0, 0, 0), UTC))
}

fn get_release_for_date(d: NaiveDate) -> DashResult<Release> {
    let url = format!("https://static.rust-lang.org/dist/{}-{:02}-{:02}/channel-rust-nightly.toml",
                      d.year(),
                      d.month(),
                      d.day());

    let response = try!(CLIENT.get(&url).send());
    match response.status {
        StatusCode::Ok => Ok(Release { date: d, released: true }),
        _ => Ok(Release { date: d, released: false }),
    }
}

pub fn get_releases_since(d: NaiveDate) -> DashResult<Vec<Release>> {
    let mut releases = vec![];

    let mut curr = d;
    let today = UTC::today().naive_utc();

    while curr <= today {
        let curr_release = try!(get_release_for_date(curr));
        releases.push(curr_release);
        curr = curr.succ();
    }

    Ok(releases)
}

pub fn ingest_releases_since(d: DateTime<UTC>) -> DashResult<()> {
    use diesel::prelude::*;
    use domain::schema::release::dsl::*;
    let releases = try!(get_releases_since(d.date().naive_utc()));

    let conn = try!(DB_POOL.get());

    for r in releases {
        let pk = release.filter(date.eq(r.date))
            .first::<(i32, NaiveDate, bool)>(&*conn)
            .map(|f| f.0)
            .ok();

        if let Some(pk) = pk {
            try!(diesel::update(release.find(pk)).set(&r).execute(&*conn));
        } else {
            try!(diesel::insert(&r).into(release).execute(&*conn));
        }
    }

    Ok(())
}
<|endoftext|>


<commit_before><commit_msg>not finished smart enum, wip for now<commit_after><|endoftext|>


<commit_before>use std;
use std::cmp::min;

use arrayvec::ArrayVec;

use coefs::Coefficients;
use params::BaseParams;
use prev::PrevFrame;

#[derive(Clone)]
pub struct Spectrals(ArrayVec<[f32; 56]>);

impl Spectrals {
    pub fn new(coefs: &Coefficients, params: &BaseParams, prev: &PrevFrame) -> Spectrals {
        let mut spectral = ArrayVec::new();

        let indexes = |x: u32| {
            let k = prev.params.harmonics as f32 / params.harmonics as f32 * x as f32;
            let int = k.floor();
            let dec = k - int;

            (int as usize, dec)
        };

        let pred = (0.03 * params.harmonics as f32 - 0.05).max(0.4).min(0.7);
        let pred_scaled = pred / params.harmonics as f32;

        let sum = (1..params.harmonics+1).map(|lambda| {
            let (k, dec) = indexes(lambda);
            (1.0 - dec) * prev.spectrals.get(k).log2() +
                dec * prev.spectrals.get(k + 1).log2()
        }).fold(0.0, |s, x| s + x);

        for l in 1..params.harmonics+1 {
            let (k, dec) = indexes(l);

            spectral.push((
                coefs.get(l as usize) +
                    pred * (1.0 - dec) * prev.spectrals.get(k).log2() +
                    pred * dec * prev.spectrals.get(k + 1).log2() -
                    pred_scaled * sum
            ).exp2());
        }

        Spectrals(spectral)
    }

    pub fn get(&self, l: usize) -> f32 {
        if l == 0 {
            1.0
        } else {
            self.0[min(l, self.0.len()) - 1]
        }
    }
}

impl std::ops::Deref for Spectrals {
    type Target = ArrayVec<[f32; 56]>;
    fn deref(&self) -> &Self::Target { &self.0 }
}

impl Default for Spectrals {
    fn default() -> Spectrals {
        Spectrals((0..56).map(|_| 1.0).collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use descramble::{Bootstrap, descramble};
    use gain::Gains;
    use coefs::Coefficients;
    use params::BaseParams;
    use prev::{PrevFrame};

    #[test]
    fn test_spectrals() {
        let chunks = [
            0b001000010010,
            0b110011001100,
            0b111000111000,
            0b111111111111,
            0b10100110101,
            0b00101111010,
            0b01110111011,
            0b00001000,
        ];

        let b = Bootstrap::new(&chunks);
        let p = BaseParams::new(b.unwrap_period());
        let (amps, _, gain_idx) = descramble(&chunks, &p);
        let g = Gains::new(gain_idx, &amps, &p);
        let c = Coefficients::new(&g, &amps, &p);
        let mut prev = PrevFrame::default();
        let s = Spectrals::new(&c, &p, &prev);

        assert_eq!(s.0.len(), 16);

        assert!((s.get(0) - 1.0).abs() < 0.000001);
        assert!((s.get(1) - 0.5306769781475001).abs() < 0.000001);
        assert!((s.get(2) - 0.3535007196676076).abs() < 0.000001);
        assert!((s.get(3) - 0.9173875577243951).abs() < 0.000001);
        assert!((s.get(4) - 0.13169278935782622).abs() < 0.000001);
        assert!((s.get(5) - 4.438599873836148).abs() < 0.000001);
        assert!((s.get(6) - 0.6796441620283217).abs() < 0.000001);
        assert!((s.get(7) - 0.9439604687782126).abs() < 0.000001);
        assert!((s.get(8) - 10.646341109175768).abs() < 0.000001);
        assert!((s.get(9) - 1.3058035207362984).abs() < 0.000001);
        assert!((s.get(10) - 1.566168983101457).abs() < 0.000001);
        assert!((s.get(11) - 8.325591838823936).abs() < 0.00001);
        assert!((s.get(12) - 1.5943520720633444).abs() < 0.000001);
        assert!((s.get(13) - 1.556318849691248).abs() < 0.000001);
        assert!((s.get(14) - 15.882295806333524).abs() < 0.00001);
        assert!((s.get(15) - 11.386537444591385).abs() < 0.00001);
        assert!((s.get(16) - 12.498792573794846).abs() < 0.00001);
        assert!((s.get(17) - 12.498792573794846).abs() < 0.00001);
        assert!((s.get(30) - 12.498792573794846).abs() < 0.00001);
        assert!((s.get(56) - 12.498792573794846).abs() < 0.00001);

        prev.spectrals = s;
        let s = Spectrals::new(&c, &p, &prev);

        assert!((s.get(1) - 0.18026009141598204).abs() < 0.000001);
        assert!((s.get(2) - 0.09466759399487303).abs() < 0.000001);
        assert!((s.get(3) - 0.546543586730317).abs() < 0.000001);
        assert!((s.get(4) - 0.11240949805284323).abs() < 0.000001);
        assert!((s.get(5) - 2.664303862945382).abs() < 0.000001);
        assert!((s.get(6) - 0.7356498197494585).abs() < 0.000001);
        assert!((s.get(7) - 0.6722918529743475).abs() < 0.000001);
        assert!((s.get(8) - 15.748045510026174).abs() < 0.000001);
        assert!((s.get(9) - 2.010522620838216).abs() < 0.000001);
        assert!((s.get(10) - 2.411402725277654).abs() < 0.000001);
        assert!((s.get(11) - 12.818766727159018).abs() < 0.00001);
        assert!((s.get(12) - 2.4547957296486485).abs() < 0.000001);
        assert!((s.get(13) - 2.396236648815911).abs() < 0.000001);
        assert!((s.get(14) - 24.45369037714974).abs() < 0.0001);
        assert!((s.get(15) - 17.53165062111628).abs() < 0.0001);
        assert!((s.get(16) - 19.244170201509185).abs() < 0.0001);
    }
}
<commit_msg>Lower tolerance of comparison<commit_after>use std;
use std::cmp::min;

use arrayvec::ArrayVec;

use coefs::Coefficients;
use params::BaseParams;
use prev::PrevFrame;

#[derive(Clone)]
pub struct Spectrals(ArrayVec<[f32; 56]>);

impl Spectrals {
    pub fn new(coefs: &Coefficients, params: &BaseParams, prev: &PrevFrame) -> Spectrals {
        let mut spectral = ArrayVec::new();

        let indexes = |x: u32| {
            let k = prev.params.harmonics as f32 / params.harmonics as f32 * x as f32;
            let int = k.floor();
            let dec = k - int;

            (int as usize, dec)
        };

        let pred = (0.03 * params.harmonics as f32 - 0.05).max(0.4).min(0.7);
        let pred_scaled = pred / params.harmonics as f32;

        let sum = (1..params.harmonics+1).map(|lambda| {
            let (k, dec) = indexes(lambda);
            (1.0 - dec) * prev.spectrals.get(k).log2() +
                dec * prev.spectrals.get(k + 1).log2()
        }).fold(0.0, |s, x| s + x);

        for l in 1..params.harmonics+1 {
            let (k, dec) = indexes(l);

            spectral.push((
                coefs.get(l as usize) +
                    pred * (1.0 - dec) * prev.spectrals.get(k).log2() +
                    pred * dec * prev.spectrals.get(k + 1).log2() -
                    pred_scaled * sum
            ).exp2());
        }

        Spectrals(spectral)
    }

    pub fn get(&self, l: usize) -> f32 {
        if l == 0 {
            1.0
        } else {
            self.0[min(l, self.0.len()) - 1]
        }
    }
}

impl std::ops::Deref for Spectrals {
    type Target = ArrayVec<[f32; 56]>;
    fn deref(&self) -> &Self::Target { &self.0 }
}

impl Default for Spectrals {
    fn default() -> Spectrals {
        Spectrals((0..56).map(|_| 1.0).collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use descramble::{Bootstrap, descramble};
    use gain::Gains;
    use coefs::Coefficients;
    use params::BaseParams;
    use prev::{PrevFrame};

    #[test]
    fn test_spectrals() {
        let chunks = [
            0b001000010010,
            0b110011001100,
            0b111000111000,
            0b111111111111,
            0b10100110101,
            0b00101111010,
            0b01110111011,
            0b00001000,
        ];

        let b = Bootstrap::new(&chunks);
        let p = BaseParams::new(b.unwrap_period());
        let (amps, _, gain_idx) = descramble(&chunks, &p);
        let g = Gains::new(gain_idx, &amps, &p);
        let c = Coefficients::new(&g, &amps, &p);
        let mut prev = PrevFrame::default();
        let s = Spectrals::new(&c, &p, &prev);

        assert_eq!(s.0.len(), 16);

        assert!((s.get(0) - 1.0).abs() < 0.000001);
        assert!((s.get(1) - 0.5306769781475001).abs() < 0.000001);
        assert!((s.get(2) - 0.3535007196676076).abs() < 0.000001);
        assert!((s.get(3) - 0.9173875577243951).abs() < 0.000001);
        assert!((s.get(4) - 0.13169278935782622).abs() < 0.000001);
        assert!((s.get(5) - 4.438599873836148).abs() < 0.00001);
        assert!((s.get(6) - 0.6796441620283217).abs() < 0.000001);
        assert!((s.get(7) - 0.9439604687782126).abs() < 0.000001);
        assert!((s.get(8) - 10.646341109175768).abs() < 0.000001);
        assert!((s.get(9) - 1.3058035207362984).abs() < 0.000001);
        assert!((s.get(10) - 1.566168983101457).abs() < 0.000001);
        assert!((s.get(11) - 8.325591838823936).abs() < 0.00001);
        assert!((s.get(12) - 1.5943520720633444).abs() < 0.000001);
        assert!((s.get(13) - 1.556318849691248).abs() < 0.000001);
        assert!((s.get(14) - 15.882295806333524).abs() < 0.00001);
        assert!((s.get(15) - 11.386537444591385).abs() < 0.00001);
        assert!((s.get(16) - 12.498792573794846).abs() < 0.00001);
        assert!((s.get(17) - 12.498792573794846).abs() < 0.00001);
        assert!((s.get(30) - 12.498792573794846).abs() < 0.00001);
        assert!((s.get(56) - 12.498792573794846).abs() < 0.00001);

        prev.spectrals = s;
        let s = Spectrals::new(&c, &p, &prev);

        assert!((s.get(1) - 0.18026009141598204).abs() < 0.000001);
        assert!((s.get(2) - 0.09466759399487303).abs() < 0.000001);
        assert!((s.get(3) - 0.546543586730317).abs() < 0.000001);
        assert!((s.get(4) - 0.11240949805284323).abs() < 0.000001);
        assert!((s.get(5) - 2.664303862945382).abs() < 0.000001);
        assert!((s.get(6) - 0.7356498197494585).abs() < 0.000001);
        assert!((s.get(7) - 0.6722918529743475).abs() < 0.000001);
        assert!((s.get(8) - 15.748045510026174).abs() < 0.000001);
        assert!((s.get(9) - 2.010522620838216).abs() < 0.000001);
        assert!((s.get(10) - 2.411402725277654).abs() < 0.000001);
        assert!((s.get(11) - 12.818766727159018).abs() < 0.00001);
        assert!((s.get(12) - 2.4547957296486485).abs() < 0.000001);
        assert!((s.get(13) - 2.396236648815911).abs() < 0.000001);
        assert!((s.get(14) - 24.45369037714974).abs() < 0.0001);
        assert!((s.get(15) - 17.53165062111628).abs() < 0.0001);
        assert!((s.get(16) - 19.244170201509185).abs() < 0.0001);
    }
}
<|endoftext|>


<commit_before>use core::cell::UnsafeCell;
use core::ops::{Deref, DerefMut};
use core::ptr;
use core::sync::atomic::{AtomicBool, AtomicPtr, Ordering, fence};

use pause::pause;

pub struct Slot {
    next: AtomicPtr<AtomicBool>
}

/// An RAII implementation of a "scoped lock" of a mutex. When this structure is
/// dropped (falls out of scope), the lock will be unlocked.
///
/// The data protected by the mutex can be access through this guard via its
/// `Deref` and `DerefMut` implementations
#[must_use]
pub struct Guard<'a, T: ?Sized + 'a> {
    lock: &'a Mutex<T>,
    slot: &'a Slot
}

/// A mutual exclusion primitive useful for protecting shared data
///
/// This mutex will block threads waiting for the lock to become available. The
/// mutex can also be statically initialized or created via a `new`
/// constructor. Each mutex has a type parameter which represents the data that
/// it is protecting. The data can only be accessed through the RAII guards
/// returned from `lock` and `try_lock`, which guarantees that the data is only
/// ever accessed when the mutex is locked.
///
/// # Examples
///
/// ```
/// use std::sync::Arc;
/// use std::thread;
/// use std::sync::mpsc::channel;
/// use mcs::{Mutex, Slot};
///
/// const N: usize = 10;
///
/// // Spawn a few threads to increment a shared variable (non-atomically), and
/// // let the main thread know once all increments are done.
/// //
/// // Here we're using an Arc to share memory among threads, and the data inside
/// // the Arc is protected with a mutex.
/// let data = Arc::new(Mutex::new(0));
///
/// let (tx, rx) = channel();
/// for _ in 0..10 {
///     let (data, tx) = (data.clone(), tx.clone());
///     thread::spawn(move || {
///         let mut slot = Slot::new();
///
///         // The shared state can only be accessed once the lock is held.
///         // Our non-atomic increment is safe because we're the only thread
///         // which can access the shared state when the lock is held.
///         //
///         // We unwrap() the return value to assert that we are not expecting
///         // threads to ever fail while holding the lock.
///         let mut data = data.lock(&mut slot);
///         *data += 1;
///         if *data == N {
///             tx.send(()).unwrap();
///         }
///         // the lock is unlocked here when `data` goes out of scope.
///     });
/// }
///
/// rx.recv().unwrap();
/// ```
pub struct Mutex<T: ?Sized> {
    queue: AtomicPtr<Slot>,
    data: UnsafeCell<T>
}

unsafe impl<T: Send> Sync for Mutex<T> { }
unsafe impl<T: Send> Send for Mutex<T> { }

impl Slot {
    pub const fn new() -> Slot {
        Slot {
            next: AtomicPtr::new(ptr::null_mut())
        }
    }
}

impl<T> Mutex<T> {
    /// Creates a new mutex in an unlocked state ready for use.
    pub const fn new(value: T) -> Mutex<T> {
        Mutex {
            queue: AtomicPtr::new(ptr::null_mut()),
            data: UnsafeCell::new(value)
        }
    }

    /// Consumes this mutex, returning the underlying data.
    pub fn into_inner(self) -> T {
        unsafe {
            self.data.into_inner()
        }
    }
}

impl<T: ?Sized> Mutex<T> {
    /// Attempts to acquire this lock.
    ///
    /// If the lock could not be acquired at this time, then `Err` is returned.
    /// Otherwise, an RAII guard is returned. The lock will be unlocked when the
    /// guard is dropped.
    ///
    /// This function does not block.
    pub fn try_lock<'a>(&'a self, slot: &'a mut Slot) -> Result<Guard<'a, T>, ()> {
        slot.next = AtomicPtr::new(ptr::null_mut());

        if self.queue.compare_and_swap(ptr::null_mut(), slot, Ordering::AcqRel).is_null() {
            Ok(Guard {
                lock: self,
                slot: slot
            })
        } else {
            Err(())
        }
    }

    /// Acquires a mutex, blocking the current thread until it is able to do so.
    ///
    /// This function will block the local thread until it is available to acquire
    /// the mutex. Upon returning, the thread is the only thread with the mutex
    /// held. An RAII guard is returned to allow scoped unlock of the lock. When
    /// the guard goes out of scope, the mutex will be unlocked.
    pub fn lock<'a>(&'a self, slot: &'a mut Slot) -> Guard<'a, T> {
        slot.next = AtomicPtr::new(ptr::null_mut());
        let pred = self.queue.swap(slot, Ordering::AcqRel);
        if !pred.is_null() {
            let pred = unsafe { &*pred };
            let locked = AtomicBool::new(true);
            pred.next.store(&locked as *const _ as *mut _, Ordering::Release);
            while locked.load(Ordering::Relaxed) {
                pause();
            }
            fence(Ordering::Acquire);
        }

        Guard {
            lock: self,
            slot: slot
        }
    }

    /// Returns a mutable reference to the underlying data.
    ///
    /// Since this call borrows the `Mutex` mutably, no actual locking needs to
    /// take place---the mutable borrow statically guarantees no locks exist.
    pub fn get_mut(&mut self) -> &mut T {
        unsafe { &mut *self.data.get() }
    }
}

impl<'a, T: ?Sized> Deref for Guard<'a, T> {
    type Target = T;
    fn deref(&self) -> &T {
        unsafe { &*self.lock.data.get() }
    }
}

impl<'a, T: ?Sized> DerefMut for Guard<'a, T> {
    fn deref_mut(&mut self) -> &mut T {
        unsafe { &mut *self.lock.data.get() }
    }
}

impl<'a, T: ?Sized> Drop for Guard<'a, T> {
    #[unsafe_destructor_blind_to_params]
    fn drop(&mut self) {
        let mut succ = self.slot.next.load(Ordering::Relaxed);
        if succ.is_null() {
            // No one has registered as waiting.
            let slot_ptr = self.slot as *const _ as *mut _;
            if self.lock.queue.compare_and_swap(slot_ptr, ptr::null_mut(), Ordering::Release) == slot_ptr {
                // No one was waiting.
                return;
            }

            // Some thread is waiting, but hasn't registered yet. Spin waiting for them to register themselves.
            loop {
                succ = self.slot.next.load(Ordering::Relaxed);
                if !succ.is_null() {
                    break;
                }
                pause();
            }
        }

        // Announce to the next waiter that the lock is free.
        fence(Ordering::Acquire);
        let succ = unsafe { &*succ };
        succ.store(false, Ordering::Release);
    }
}

#[cfg(test)]
mod test {
    extern crate std;

    use super::{Mutex, Slot};

    // Mostly stoled from the Rust standard Mutex implementation's tests, so

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

    use self::std::sync::Arc;
    use self::std::sync::mpsc::channel;
    use self::std::sync::atomic::{AtomicUsize, Ordering};
    use self::std::thread;

    #[derive(Eq, PartialEq, Debug)]
    struct NonCopy(i32);

    #[test]
    fn smoke() {
        let mut slot = Slot::new();
        let m = Mutex::new(());
        drop(m.lock(&mut slot));
        drop(m.lock(&mut slot));
    }

    #[test]
    fn lots_and_lots() {
        static LOCK: Mutex<u32> = Mutex::new(0);
        const ITERS: u32 = 1000;
        const CONCURRENCY: u32 = 3;

        fn inc() {
            let mut slot = Slot::new();
            for _ in 0..ITERS {
                let mut g = LOCK.lock(&mut slot);
                *g += 1;
            }
        };

        let (tx, rx) = channel();
        for _ in 0..CONCURRENCY {
            let tx2 = tx.clone();
            thread::spawn(move|| { inc(); tx2.send(()).unwrap(); });
            let tx2 = tx.clone();
            thread::spawn(move|| { inc(); tx2.send(()).unwrap(); });
        }

        drop(tx);
        for _ in 0..2 * CONCURRENCY {
            rx.recv().unwrap();
        }
        let mut slot = Slot::new();
        assert_eq!(*LOCK.lock(&mut slot), ITERS * CONCURRENCY * 2);
    }

    #[test]
    fn try_lock() {
        let mut slot = Slot::new();
        let m = Mutex::new(());
        *m.try_lock(&mut slot).unwrap() = ();
    }

    #[test]
    fn test_into_inner() {
        let m = Mutex::new(NonCopy(10));
        assert_eq!(m.into_inner(), NonCopy(10));
    }

    #[test]
    fn test_into_inner_drop() {
        struct Foo(Arc<AtomicUsize>);
        impl Drop for Foo {
            fn drop(&mut self) {
                self.0.fetch_add(1, Ordering::SeqCst);
            }
        }
        let num_drops = Arc::new(AtomicUsize::new(0));
        let m = Mutex::new(Foo(num_drops.clone()));
        assert_eq!(num_drops.load(Ordering::SeqCst), 0);
        {
            let _inner = m.into_inner();
            assert_eq!(num_drops.load(Ordering::SeqCst), 0);
        }
        assert_eq!(num_drops.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_get_mut() {
        let mut m = Mutex::new(NonCopy(10));
        *m.get_mut() = NonCopy(20);
        assert_eq!(m.into_inner(), NonCopy(20));
    }

    #[test]
    fn test_lock_arc_nested() {
        // Tests nested locks and access
        // to underlying data.
        let arc = Arc::new(Mutex::new(1));
        let arc2 = Arc::new(Mutex::new(arc));
        let (tx, rx) = channel();
        let _t = thread::spawn(move|| {
            let mut slot1 = Slot::new();
            let mut slot2 = Slot::new();

            let lock = arc2.lock(&mut slot1);
            let lock2 = lock.lock(&mut slot2);
            assert_eq!(*lock2, 1);
            tx.send(()).unwrap();
        });
        rx.recv().unwrap();
    }

    #[test]
    fn test_lock_arc_access_in_unwind() {
        let arc = Arc::new(Mutex::new(1));
        let arc2 = arc.clone();
        let _ = thread::spawn(move|| -> () {
            struct Unwinder {
                i: Arc<Mutex<i32>>,
            }
            impl Drop for Unwinder {
                fn drop(&mut self) {
                    let mut slot = Slot::new();
                    *self.i.lock(&mut slot) += 1;
                }
            }
            let _u = Unwinder { i: arc2 };
            panic!();
        }).join();
        let mut slot = Slot::new();
        let lock = arc.lock(&mut slot);
        assert_eq!(*lock, 2);
    }

    #[test]
    fn test_lock_unsized() {
        let mut slot = Slot::new();
        let lock: &Mutex<[i32]> = &Mutex::new([1, 2, 3]);
        {
            let b = &mut *lock.lock(&mut slot);
            b[0] = 4;
            b[2] = 5;
        }
        let comp: &[i32] = &[4, 2, 5];
        assert_eq!(&*lock.lock(&mut slot), comp);
    }
}
<commit_msg>Use compare_exchange in Guard's Drop impl<commit_after>use core::cell::UnsafeCell;
use core::ops::{Deref, DerefMut};
use core::ptr;
use core::sync::atomic::{AtomicBool, AtomicPtr, Ordering, fence};

use pause::pause;

pub struct Slot {
    next: AtomicPtr<AtomicBool>
}

/// An RAII implementation of a "scoped lock" of a mutex. When this structure is
/// dropped (falls out of scope), the lock will be unlocked.
///
/// The data protected by the mutex can be access through this guard via its
/// `Deref` and `DerefMut` implementations
#[must_use]
pub struct Guard<'a, T: ?Sized + 'a> {
    lock: &'a Mutex<T>,
    slot: &'a Slot
}

/// A mutual exclusion primitive useful for protecting shared data
///
/// This mutex will block threads waiting for the lock to become available. The
/// mutex can also be statically initialized or created via a `new`
/// constructor. Each mutex has a type parameter which represents the data that
/// it is protecting. The data can only be accessed through the RAII guards
/// returned from `lock` and `try_lock`, which guarantees that the data is only
/// ever accessed when the mutex is locked.
///
/// # Examples
///
/// ```
/// use std::sync::Arc;
/// use std::thread;
/// use std::sync::mpsc::channel;
/// use mcs::{Mutex, Slot};
///
/// const N: usize = 10;
///
/// // Spawn a few threads to increment a shared variable (non-atomically), and
/// // let the main thread know once all increments are done.
/// //
/// // Here we're using an Arc to share memory among threads, and the data inside
/// // the Arc is protected with a mutex.
/// let data = Arc::new(Mutex::new(0));
///
/// let (tx, rx) = channel();
/// for _ in 0..10 {
///     let (data, tx) = (data.clone(), tx.clone());
///     thread::spawn(move || {
///         let mut slot = Slot::new();
///
///         // The shared state can only be accessed once the lock is held.
///         // Our non-atomic increment is safe because we're the only thread
///         // which can access the shared state when the lock is held.
///         //
///         // We unwrap() the return value to assert that we are not expecting
///         // threads to ever fail while holding the lock.
///         let mut data = data.lock(&mut slot);
///         *data += 1;
///         if *data == N {
///             tx.send(()).unwrap();
///         }
///         // the lock is unlocked here when `data` goes out of scope.
///     });
/// }
///
/// rx.recv().unwrap();
/// ```
pub struct Mutex<T: ?Sized> {
    queue: AtomicPtr<Slot>,
    data: UnsafeCell<T>
}

unsafe impl<T: Send> Sync for Mutex<T> { }
unsafe impl<T: Send> Send for Mutex<T> { }

impl Slot {
    pub const fn new() -> Slot {
        Slot {
            next: AtomicPtr::new(ptr::null_mut())
        }
    }
}

impl<T> Mutex<T> {
    /// Creates a new mutex in an unlocked state ready for use.
    pub const fn new(value: T) -> Mutex<T> {
        Mutex {
            queue: AtomicPtr::new(ptr::null_mut()),
            data: UnsafeCell::new(value)
        }
    }

    /// Consumes this mutex, returning the underlying data.
    pub fn into_inner(self) -> T {
        unsafe {
            self.data.into_inner()
        }
    }
}

impl<T: ?Sized> Mutex<T> {
    /// Attempts to acquire this lock.
    ///
    /// If the lock could not be acquired at this time, then `Err` is returned.
    /// Otherwise, an RAII guard is returned. The lock will be unlocked when the
    /// guard is dropped.
    ///
    /// This function does not block.
    pub fn try_lock<'a>(&'a self, slot: &'a mut Slot) -> Result<Guard<'a, T>, ()> {
        slot.next = AtomicPtr::new(ptr::null_mut());

        if self.queue.compare_and_swap(ptr::null_mut(), slot, Ordering::AcqRel).is_null() {
            Ok(Guard {
                lock: self,
                slot: slot
            })
        } else {
            Err(())
        }
    }

    /// Acquires a mutex, blocking the current thread until it is able to do so.
    ///
    /// This function will block the local thread until it is available to acquire
    /// the mutex. Upon returning, the thread is the only thread with the mutex
    /// held. An RAII guard is returned to allow scoped unlock of the lock. When
    /// the guard goes out of scope, the mutex will be unlocked.
    pub fn lock<'a>(&'a self, slot: &'a mut Slot) -> Guard<'a, T> {
        slot.next = AtomicPtr::new(ptr::null_mut());
        let pred = self.queue.swap(slot, Ordering::AcqRel);
        if !pred.is_null() {
            let pred = unsafe { &*pred };
            let locked = AtomicBool::new(true);
            pred.next.store(&locked as *const _ as *mut _, Ordering::Release);
            while locked.load(Ordering::Relaxed) {
                pause();
            }
            fence(Ordering::Acquire);
        }

        Guard {
            lock: self,
            slot: slot
        }
    }

    /// Returns a mutable reference to the underlying data.
    ///
    /// Since this call borrows the `Mutex` mutably, no actual locking needs to
    /// take place---the mutable borrow statically guarantees no locks exist.
    pub fn get_mut(&mut self) -> &mut T {
        unsafe { &mut *self.data.get() }
    }
}

impl<'a, T: ?Sized> Deref for Guard<'a, T> {
    type Target = T;
    fn deref(&self) -> &T {
        unsafe { &*self.lock.data.get() }
    }
}

impl<'a, T: ?Sized> DerefMut for Guard<'a, T> {
    fn deref_mut(&mut self) -> &mut T {
        unsafe { &mut *self.lock.data.get() }
    }
}

impl<'a, T: ?Sized> Drop for Guard<'a, T> {
    #[unsafe_destructor_blind_to_params]
    fn drop(&mut self) {
        let mut succ = self.slot.next.load(Ordering::Relaxed);
        if succ.is_null() {
            // No one has registered as waiting.
            if self.lock.queue.compare_exchange(self.slot as *const _ as *mut _, ptr::null_mut(), Ordering::Release, Ordering::Relaxed).is_ok() {
                // No one was waiting.
                return;
            }

            // Some thread is waiting, but hasn't registered yet. Spin waiting for them to register themselves.
            loop {
                succ = self.slot.next.load(Ordering::Relaxed);
                if !succ.is_null() {
                    break;
                }
                pause();
            }
        }

        // Announce to the next waiter that the lock is free.
        fence(Ordering::Acquire);
        let succ = unsafe { &*succ };
        succ.store(false, Ordering::Release);
    }
}

#[cfg(test)]
mod test {
    extern crate std;

    use super::{Mutex, Slot};

    // Mostly stoled from the Rust standard Mutex implementation's tests, so

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

    use self::std::sync::Arc;
    use self::std::sync::mpsc::channel;
    use self::std::sync::atomic::{AtomicUsize, Ordering};
    use self::std::thread;

    #[derive(Eq, PartialEq, Debug)]
    struct NonCopy(i32);

    #[test]
    fn smoke() {
        let mut slot = Slot::new();
        let m = Mutex::new(());
        drop(m.lock(&mut slot));
        drop(m.lock(&mut slot));
    }

    #[test]
    fn lots_and_lots() {
        static LOCK: Mutex<u32> = Mutex::new(0);
        const ITERS: u32 = 1000;
        const CONCURRENCY: u32 = 3;

        fn inc() {
            let mut slot = Slot::new();
            for _ in 0..ITERS {
                let mut g = LOCK.lock(&mut slot);
                *g += 1;
            }
        };

        let (tx, rx) = channel();
        for _ in 0..CONCURRENCY {
            let tx2 = tx.clone();
            thread::spawn(move|| { inc(); tx2.send(()).unwrap(); });
            let tx2 = tx.clone();
            thread::spawn(move|| { inc(); tx2.send(()).unwrap(); });
        }

        drop(tx);
        for _ in 0..2 * CONCURRENCY {
            rx.recv().unwrap();
        }
        let mut slot = Slot::new();
        assert_eq!(*LOCK.lock(&mut slot), ITERS * CONCURRENCY * 2);
    }

    #[test]
    fn try_lock() {
        let mut slot = Slot::new();
        let m = Mutex::new(());
        *m.try_lock(&mut slot).unwrap() = ();
    }

    #[test]
    fn test_into_inner() {
        let m = Mutex::new(NonCopy(10));
        assert_eq!(m.into_inner(), NonCopy(10));
    }

    #[test]
    fn test_into_inner_drop() {
        struct Foo(Arc<AtomicUsize>);
        impl Drop for Foo {
            fn drop(&mut self) {
                self.0.fetch_add(1, Ordering::SeqCst);
            }
        }
        let num_drops = Arc::new(AtomicUsize::new(0));
        let m = Mutex::new(Foo(num_drops.clone()));
        assert_eq!(num_drops.load(Ordering::SeqCst), 0);
        {
            let _inner = m.into_inner();
            assert_eq!(num_drops.load(Ordering::SeqCst), 0);
        }
        assert_eq!(num_drops.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_get_mut() {
        let mut m = Mutex::new(NonCopy(10));
        *m.get_mut() = NonCopy(20);
        assert_eq!(m.into_inner(), NonCopy(20));
    }

    #[test]
    fn test_lock_arc_nested() {
        // Tests nested locks and access
        // to underlying data.
        let arc = Arc::new(Mutex::new(1));
        let arc2 = Arc::new(Mutex::new(arc));
        let (tx, rx) = channel();
        let _t = thread::spawn(move|| {
            let mut slot1 = Slot::new();
            let mut slot2 = Slot::new();

            let lock = arc2.lock(&mut slot1);
            let lock2 = lock.lock(&mut slot2);
            assert_eq!(*lock2, 1);
            tx.send(()).unwrap();
        });
        rx.recv().unwrap();
    }

    #[test]
    fn test_lock_arc_access_in_unwind() {
        let arc = Arc::new(Mutex::new(1));
        let arc2 = arc.clone();
        let _ = thread::spawn(move|| -> () {
            struct Unwinder {
                i: Arc<Mutex<i32>>,
            }
            impl Drop for Unwinder {
                fn drop(&mut self) {
                    let mut slot = Slot::new();
                    *self.i.lock(&mut slot) += 1;
                }
            }
            let _u = Unwinder { i: arc2 };
            panic!();
        }).join();
        let mut slot = Slot::new();
        let lock = arc.lock(&mut slot);
        assert_eq!(*lock, 2);
    }

    #[test]
    fn test_lock_unsized() {
        let mut slot = Slot::new();
        let lock: &Mutex<[i32]> = &Mutex::new([1, 2, 3]);
        {
            let b = &mut *lock.lock(&mut slot);
            b[0] = 4;
            b[2] = 5;
        }
        let comp: &[i32] = &[4, 2, 5];
        assert_eq!(&*lock.lock(&mut slot), comp);
    }
}
<|endoftext|>


<commit_before>//! Functions for parsing TIS-100 assembly code into instructions.

use std::str::FromStr;
use std::fmt::{Display, Formatter, Error};
use std::collections::HashMap;
use core::*;
use core::Instruction::*;
use lex::{lex_program, Label, Line};

/// An error that can be returned while parsing a TIS-100 assembly program.
#[derive(Debug, PartialEq)]
pub enum ParseProgramError {
    InvalidLabel,
    UndefinedLabel(String),
    DuplicateLabel(String),
    InvalidOpcode(String),
    InvalidExpression(String),
    InvalidRegister(String),
    MissingOperand(String),
    TooManyOperands(String),
}

impl Display for ParseProgramError {
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        match self {
            &InvalidLabel => f.write_str("Invalid label"),
            &UndefinedLabel(ref lbl) => f.write_fmt(format_args!("Undefined label: '{}'", lbl)),
            &DuplicateLabel(ref lbl) => f.write_fmt(format_args!("Label is already defined: '{}'", lbl)),
            &InvalidOpcode(ref op) => f.write_fmt(format_args!("Invalid opcode: '{}'", op)),
            &InvalidExpression(ref expr) => f.write_fmt(format_args!("Invalid expression: '{}'", expr)),
            &InvalidRegister(ref reg) => f.write_fmt(format_args!("Invalid register: '{}'", reg)),
            &MissingOperand(ref op) => f.write_fmt(format_args!("Missing operand: '{}'", op)),
            &TooManyOperands(ref ops) => f.write_fmt(format_args!("Too many operands: '{}'", ops)),
        }
    }
}

/// All errors discovered while parsing a TIS-100 assembly program.
pub type ProgramErrors = Vec<(usize, ParseProgramError)>;

use self::ParseProgramError::*;

/// A result that can be returned intermediate phases of the parsing process.
type ParseResult<T> = Result<T, ParseProgramError>;

/// Parse the program source code into a list of instructions. If one or more errors are
/// encountered during parsing, then the list of errors will be returned instead.
///
/// # Example
///
/// ```
/// use tis_100::core::Instruction::*;
/// use tis_100::core::Source::*;
/// use tis_100::core::Register::*;
/// use tis_100::core::IoRegister::*;
/// use tis_100::core::Port::*;
/// use tis_100::parse::parse_program;
///
/// let src = "MOV UP ACC\nADD 1\nMOV ACC DOWN\n";
/// let prog = parse_program(src).unwrap();
/// assert_eq!(prog[0], Mov(REG(IO(DIR(UP))), ACC));
/// assert_eq!(prog[1], Add(VAL(1)));
/// assert_eq!(prog[2], Mov(REG(ACC), IO(DIR(DOWN))));
/// ```
pub fn parse_program(src: &str) -> Result<Program, ProgramErrors> {
    // The basic parsing process is:
    // 1. Tokenize the source into labels, opcodes, and operands
    // 2. Create a mapping of labels to instruction indices
    // 3. Parse opcodes and operands line-by-line to generate instructions

    let mut label_map = HashMap::new();
    let mut instructions = Vec::new();
    let mut errors = Vec::new();

    let lines = lex_program(src);

    // Lable mapping pass
    for &Line(line_num, ref maybe_label, _) in lines.iter() {
        if let &Some(Label(ref name, index)) = maybe_label {
            if name.len() == 0 {
                errors.push((line_num, InvalidLabel));
            } else if let None = label_map.get(name) {
                label_map.insert(name.clone(), index as isize);
            } else {
                errors.push((line_num, DuplicateLabel(name.clone())));
            }
        }
    }

    // Instruction pass
    for &Line(line_num, _, ref lexemes) in lines.iter() {
        if lexemes.len() > 0 {
            match parse_instruction(&lexemes[0], &lexemes[1..], &label_map) {
                Ok(instruction) => instructions.push(instruction),
                Err(err) => errors.push((line_num, err)),
            }
        }
    }

    if errors.len() > 0 {
        Err(errors)
    } else {
        Ok(instructions)
    }
}

/// Attempt to parse a single TIS-100 assembly instruction.
fn parse_instruction(opcode: &str, operands: &[String], labels: &HashMap<String, isize>) -> ParseResult<Instruction> {
    match str::parse::<Opcode>(opcode) {
        Ok(NOP) => parse_no_operands(Nop, operands),
        Ok(MOV) => parse_two_operands(Mov, opcode, operands),
        Ok(SWP) => parse_no_operands(Swp, operands),
        Ok(SAV) => parse_no_operands(Sav, operands),
        Ok(ADD) => parse_one_operand(Add, opcode, operands),
        Ok(SUB) => parse_one_operand(Sub, opcode, operands),
        Ok(NEG) => parse_no_operands(Neg, operands),
        Ok(JMP) => parse_jump(Jmp, opcode, operands, labels),
        Ok(JEZ) => parse_jump(Jez, opcode, operands, labels),
        Ok(JNZ) => parse_jump(Jnz, opcode, operands, labels),
        Ok(JGZ) => parse_jump(Jgz, opcode, operands, labels),
        Ok(JLZ) => parse_jump(Jlz, opcode, operands, labels),
        Ok(JRO) => parse_one_operand(Jro, opcode, operands),
        _ => Err(InvalidOpcode(opcode.to_string())),
    }
}

/// Attempt to resolve a jump label to an instruction index.
fn resolve_label<'a>(label: &str, labels: &'a HashMap<String, isize>) -> ParseResult<&'a isize> {
    labels.get(label).ok_or(UndefinedLabel(label.to_string()))
}

/// Parse a jump opcode and label into a jump instruction.
fn parse_jump<F: Fn(isize) -> Instruction>(f: F, opcode: &str, operands: &[String], labels: &HashMap<String, isize>) -> ParseResult<Instruction> {
    if operands.len() < 1 {
        Err(MissingOperand(opcode.to_string()))
    } else if operands.len() == 1 {
        resolve_label(&operands[0], labels).map(|&i| f(i))
    } else {
        Err(TooManyOperands(operands[1..].connect(" ")))
    }
}

/// Parse an opcode into an instruction.
fn parse_no_operands(instruction: Instruction, operands: &[String]) -> ParseResult<Instruction> {
    if operands.len() == 0 {
        Ok(instruction)
    } else {
        Err(TooManyOperands(operands.connect(" ")))
    }
}

/// Parse an opcode and one operand into an instruction.
fn parse_one_operand<T: FromStr, F: Fn(T) -> Instruction>(f: F, opcode: &str, operands: &[String]) -> ParseResult<Instruction> {
    if operands.len() < 1 {
        Err(MissingOperand(opcode.to_string()))
    } else if operands.len() == 1 {
        match str::parse::<T>(&operands[0]) {
            Ok(op) => Ok(f(op)),
            Err(_) => Err(InvalidExpression(operands[0].clone())),
        }
    } else {
        Err(TooManyOperands(operands[1..].connect(" ")))
    }
}

/// Parse an opcode and two operands into an instruction.
fn parse_two_operands<T: FromStr, U: FromStr, F: Fn(T, U) -> Instruction>(f: F, opcode: &str, operands: &[String]) -> ParseResult<Instruction> {
    if operands.len() < 2 {
        Err(MissingOperand(opcode.to_string() + " " + &operands.connect(" ")))
    } else if operands.len() == 2 {
        match str::parse::<T>(&operands[0]) {
            Ok(op1) => match str::parse::<U>(&operands[1]) {
                Ok(op2) => Ok(f(op1, op2)),
                Err(_) => Err(InvalidRegister(operands[1].clone())),
            },
            Err(_) => Err(InvalidExpression(operands[0].clone())),
        }
    } else {
        Err(TooManyOperands(operands[2..].connect(" ")))
    }
}

/// The opcode component of a TIS-100 instruction.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Opcode {
    NOP,
    MOV,
    SWP,
    SAV,
    ADD,
    SUB,
    NEG,
    JMP,
    JEZ,
    JNZ,
    JGZ,
    JLZ,
    JRO,
}

use self::Opcode::*;

/// An error which can be returned when parsing an opcode.
#[derive(Debug, PartialEq)]
pub struct ParseOpcodeError;

impl FromStr for Opcode {
    type Err = ParseOpcodeError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "NOP" => Ok(NOP),
            "MOV" => Ok(MOV),
            "SWP" => Ok(SWP),
            "SAV" => Ok(SAV),
            "ADD" => Ok(ADD),
            "SUB" => Ok(SUB),
            "NEG" => Ok(NEG),
            "JMP" => Ok(JMP),
            "JEZ" => Ok(JEZ),
            "JNZ" => Ok(JNZ),
            "JGZ" => Ok(JGZ),
            "JLZ" => Ok(JLZ),
            "JRO" => Ok(JRO),
            _ => Err(ParseOpcodeError),
        }
    }
}

#[test]
fn test_parse_opcode() {
    assert_eq!(str::parse::<Opcode>("NOP"), Ok(NOP));
    assert_eq!(str::parse::<Opcode>("MOV"), Ok(MOV));
    assert_eq!(str::parse::<Opcode>("SWP"), Ok(SWP));
    assert_eq!(str::parse::<Opcode>("SAV"), Ok(SAV));
    assert_eq!(str::parse::<Opcode>("ADD"), Ok(ADD));
    assert_eq!(str::parse::<Opcode>("SUB"), Ok(SUB));
    assert_eq!(str::parse::<Opcode>("NEG"), Ok(NEG));
    assert_eq!(str::parse::<Opcode>("JMP"), Ok(JMP));
    assert_eq!(str::parse::<Opcode>("JEZ"), Ok(JEZ));
    assert_eq!(str::parse::<Opcode>("JNZ"), Ok(JNZ));
    assert_eq!(str::parse::<Opcode>("JGZ"), Ok(JGZ));
    assert_eq!(str::parse::<Opcode>("JLZ"), Ok(JLZ));
    assert_eq!(str::parse::<Opcode>("JRO"), Ok(JRO));
    assert_eq!(str::parse::<Opcode>("nop"), Err(ParseOpcodeError));
    assert_eq!(str::parse::<Opcode>("bad"), Err(ParseOpcodeError));
}
<commit_msg>Replace deprecated .connect with .join<commit_after>//! Functions for parsing TIS-100 assembly code into instructions.

use std::str::FromStr;
use std::fmt::{Display, Formatter, Error};
use std::collections::HashMap;
use core::*;
use core::Instruction::*;
use lex::{lex_program, Label, Line};

/// An error that can be returned while parsing a TIS-100 assembly program.
#[derive(Debug, PartialEq)]
pub enum ParseProgramError {
    InvalidLabel,
    UndefinedLabel(String),
    DuplicateLabel(String),
    InvalidOpcode(String),
    InvalidExpression(String),
    InvalidRegister(String),
    MissingOperand(String),
    TooManyOperands(String),
}

impl Display for ParseProgramError {
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        match self {
            &InvalidLabel => f.write_str("Invalid label"),
            &UndefinedLabel(ref lbl) => f.write_fmt(format_args!("Undefined label: '{}'", lbl)),
            &DuplicateLabel(ref lbl) => f.write_fmt(format_args!("Label is already defined: '{}'", lbl)),
            &InvalidOpcode(ref op) => f.write_fmt(format_args!("Invalid opcode: '{}'", op)),
            &InvalidExpression(ref expr) => f.write_fmt(format_args!("Invalid expression: '{}'", expr)),
            &InvalidRegister(ref reg) => f.write_fmt(format_args!("Invalid register: '{}'", reg)),
            &MissingOperand(ref op) => f.write_fmt(format_args!("Missing operand: '{}'", op)),
            &TooManyOperands(ref ops) => f.write_fmt(format_args!("Too many operands: '{}'", ops)),
        }
    }
}

/// All errors discovered while parsing a TIS-100 assembly program.
pub type ProgramErrors = Vec<(usize, ParseProgramError)>;

use self::ParseProgramError::*;

/// A result that can be returned intermediate phases of the parsing process.
type ParseResult<T> = Result<T, ParseProgramError>;

/// Parse the program source code into a list of instructions. If one or more errors are
/// encountered during parsing, then the list of errors will be returned instead.
///
/// # Example
///
/// ```
/// use tis_100::core::Instruction::*;
/// use tis_100::core::Source::*;
/// use tis_100::core::Register::*;
/// use tis_100::core::IoRegister::*;
/// use tis_100::core::Port::*;
/// use tis_100::parse::parse_program;
///
/// let src = "MOV UP ACC\nADD 1\nMOV ACC DOWN\n";
/// let prog = parse_program(src).unwrap();
/// assert_eq!(prog[0], Mov(REG(IO(DIR(UP))), ACC));
/// assert_eq!(prog[1], Add(VAL(1)));
/// assert_eq!(prog[2], Mov(REG(ACC), IO(DIR(DOWN))));
/// ```
pub fn parse_program(src: &str) -> Result<Program, ProgramErrors> {
    // The basic parsing process is:
    // 1. Tokenize the source into labels, opcodes, and operands
    // 2. Create a mapping of labels to instruction indices
    // 3. Parse opcodes and operands line-by-line to generate instructions

    let mut label_map = HashMap::new();
    let mut instructions = Vec::new();
    let mut errors = Vec::new();

    let lines = lex_program(src);

    // Lable mapping pass
    for &Line(line_num, ref maybe_label, _) in lines.iter() {
        if let &Some(Label(ref name, index)) = maybe_label {
            if name.len() == 0 {
                errors.push((line_num, InvalidLabel));
            } else if let None = label_map.get(name) {
                label_map.insert(name.clone(), index as isize);
            } else {
                errors.push((line_num, DuplicateLabel(name.clone())));
            }
        }
    }

    // Instruction pass
    for &Line(line_num, _, ref lexemes) in lines.iter() {
        if lexemes.len() > 0 {
            match parse_instruction(&lexemes[0], &lexemes[1..], &label_map) {
                Ok(instruction) => instructions.push(instruction),
                Err(err) => errors.push((line_num, err)),
            }
        }
    }

    if errors.len() > 0 {
        Err(errors)
    } else {
        Ok(instructions)
    }
}

/// Attempt to parse a single TIS-100 assembly instruction.
fn parse_instruction(opcode: &str, operands: &[String], labels: &HashMap<String, isize>) -> ParseResult<Instruction> {
    match str::parse::<Opcode>(opcode) {
        Ok(NOP) => parse_no_operands(Nop, operands),
        Ok(MOV) => parse_two_operands(Mov, opcode, operands),
        Ok(SWP) => parse_no_operands(Swp, operands),
        Ok(SAV) => parse_no_operands(Sav, operands),
        Ok(ADD) => parse_one_operand(Add, opcode, operands),
        Ok(SUB) => parse_one_operand(Sub, opcode, operands),
        Ok(NEG) => parse_no_operands(Neg, operands),
        Ok(JMP) => parse_jump(Jmp, opcode, operands, labels),
        Ok(JEZ) => parse_jump(Jez, opcode, operands, labels),
        Ok(JNZ) => parse_jump(Jnz, opcode, operands, labels),
        Ok(JGZ) => parse_jump(Jgz, opcode, operands, labels),
        Ok(JLZ) => parse_jump(Jlz, opcode, operands, labels),
        Ok(JRO) => parse_one_operand(Jro, opcode, operands),
        _ => Err(InvalidOpcode(opcode.to_string())),
    }
}

/// Attempt to resolve a jump label to an instruction index.
fn resolve_label<'a>(label: &str, labels: &'a HashMap<String, isize>) -> ParseResult<&'a isize> {
    labels.get(label).ok_or(UndefinedLabel(label.to_string()))
}

/// Parse a jump opcode and label into a jump instruction.
fn parse_jump<F: Fn(isize) -> Instruction>(f: F, opcode: &str, operands: &[String], labels: &HashMap<String, isize>) -> ParseResult<Instruction> {
    if operands.len() < 1 {
        Err(MissingOperand(opcode.to_string()))
    } else if operands.len() == 1 {
        resolve_label(&operands[0], labels).map(|&i| f(i))
    } else {
        Err(TooManyOperands(operands[1..].join(" ")))
    }
}

/// Parse an opcode into an instruction.
fn parse_no_operands(instruction: Instruction, operands: &[String]) -> ParseResult<Instruction> {
    if operands.len() == 0 {
        Ok(instruction)
    } else {
        Err(TooManyOperands(operands.join(" ")))
    }
}

/// Parse an opcode and one operand into an instruction.
fn parse_one_operand<T: FromStr, F: Fn(T) -> Instruction>(f: F, opcode: &str, operands: &[String]) -> ParseResult<Instruction> {
    if operands.len() < 1 {
        Err(MissingOperand(opcode.to_string()))
    } else if operands.len() == 1 {
        match str::parse::<T>(&operands[0]) {
            Ok(op) => Ok(f(op)),
            Err(_) => Err(InvalidExpression(operands[0].clone())),
        }
    } else {
        Err(TooManyOperands(operands[1..].join(" ")))
    }
}

/// Parse an opcode and two operands into an instruction.
fn parse_two_operands<T: FromStr, U: FromStr, F: Fn(T, U) -> Instruction>(f: F, opcode: &str, operands: &[String]) -> ParseResult<Instruction> {
    if operands.len() < 2 {
        Err(MissingOperand(opcode.to_string() + " " + &operands.join(" ")))
    } else if operands.len() == 2 {
        match str::parse::<T>(&operands[0]) {
            Ok(op1) => match str::parse::<U>(&operands[1]) {
                Ok(op2) => Ok(f(op1, op2)),
                Err(_) => Err(InvalidRegister(operands[1].clone())),
            },
            Err(_) => Err(InvalidExpression(operands[0].clone())),
        }
    } else {
        Err(TooManyOperands(operands[2..].join(" ")))
    }
}

/// The opcode component of a TIS-100 instruction.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Opcode {
    NOP,
    MOV,
    SWP,
    SAV,
    ADD,
    SUB,
    NEG,
    JMP,
    JEZ,
    JNZ,
    JGZ,
    JLZ,
    JRO,
}

use self::Opcode::*;

/// An error which can be returned when parsing an opcode.
#[derive(Debug, PartialEq)]
pub struct ParseOpcodeError;

impl FromStr for Opcode {
    type Err = ParseOpcodeError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "NOP" => Ok(NOP),
            "MOV" => Ok(MOV),
            "SWP" => Ok(SWP),
            "SAV" => Ok(SAV),
            "ADD" => Ok(ADD),
            "SUB" => Ok(SUB),
            "NEG" => Ok(NEG),
            "JMP" => Ok(JMP),
            "JEZ" => Ok(JEZ),
            "JNZ" => Ok(JNZ),
            "JGZ" => Ok(JGZ),
            "JLZ" => Ok(JLZ),
            "JRO" => Ok(JRO),
            _ => Err(ParseOpcodeError),
        }
    }
}

#[test]
fn test_parse_opcode() {
    assert_eq!(str::parse::<Opcode>("NOP"), Ok(NOP));
    assert_eq!(str::parse::<Opcode>("MOV"), Ok(MOV));
    assert_eq!(str::parse::<Opcode>("SWP"), Ok(SWP));
    assert_eq!(str::parse::<Opcode>("SAV"), Ok(SAV));
    assert_eq!(str::parse::<Opcode>("ADD"), Ok(ADD));
    assert_eq!(str::parse::<Opcode>("SUB"), Ok(SUB));
    assert_eq!(str::parse::<Opcode>("NEG"), Ok(NEG));
    assert_eq!(str::parse::<Opcode>("JMP"), Ok(JMP));
    assert_eq!(str::parse::<Opcode>("JEZ"), Ok(JEZ));
    assert_eq!(str::parse::<Opcode>("JNZ"), Ok(JNZ));
    assert_eq!(str::parse::<Opcode>("JGZ"), Ok(JGZ));
    assert_eq!(str::parse::<Opcode>("JLZ"), Ok(JLZ));
    assert_eq!(str::parse::<Opcode>("JRO"), Ok(JRO));
    assert_eq!(str::parse::<Opcode>("nop"), Err(ParseOpcodeError));
    assert_eq!(str::parse::<Opcode>("bad"), Err(ParseOpcodeError));
}
<|endoftext|>