<commit_before>//! Reader and writer for the refFlat format.
//!
//! The refFlat format is a transcript-oriented format in which each
//! transcript is denoted in a single line. It is used most prominently
//! by the [picard suite tools](https://broadinstitute.github.io/picard/)
//!
//! A minimum specification of the columns can be found here:
//! https://genome.ucsc.edu/goldenPath/gbdDescriptionsOld.html#RefFlat

use std::cmp::{max, min};
use std::convert::AsRef;
use std::io;
use std::fs;
use std::path::Path;
use std::str::FromStr;

use csv;
use itertools::{GroupBy, Group, Itertools};
use linked_hash_map::LinkedHashMap;

use feature::FeatureError;
use {Coord, Gene, GBuilder, Strand, Transcript, TBuilder, Error};

pub type RefFlatRow = (String, String, String, char, u64, u64, u64, u64, usize, String, String);

pub struct RefFlatRecord {
    pub gene_id: String,
    pub transcript_name: String,
    pub seq_name: String,
    pub strand: char,
    pub trx_start: u64,
    pub trx_end: u64,
    pub coding_start: u64,
    pub coding_end: u64,
    pub num_exons: usize,
    pub exon_starts: String,
    pub exon_ends: String,
}

impl RefFlatRecord {

    pub fn into_transcript(self) -> Result<Transcript, Error> {

        let exon_coords = self.zip_raw_exon_coords()?;
        if exon_coords.len() != self.num_exons {
            return Err(Error::RefFlat(
                "number of exon and exon coordinates mismatch"));
        }

        let strand = Strand::from_char(&self.strand)
            .map_err(FeatureError::from)
            .map_err(Error::from)?;

        let coding_coord =
            if self.coding_start != self.coding_end {
                Some((self.coding_start, self.coding_end))
            } else {
                None
            };

        TBuilder::new(self.seq_name, self.trx_start, self.trx_end)
            .id(self.transcript_name)
            .gene_id(self.gene_id)
            .strand(strand)
            .coords(exon_coords, coding_coord)
            .coding_incl_stop(true)
            .build()
            .map_err(Error::from)
    }

    #[inline]
    fn zip_raw_exon_coords(&self) -> Result<Vec<Coord<u64>>, Error> {
        let exon_starts = self.exon_starts
            .trim_matches(',').split(',').map(|item| u64::from_str(item).map_err(Error::from));
        let exon_ends = self.exon_ends
            .trim_matches(',').split(',').map(|item| u64::from_str(item).map_err(Error::from));

        let mut res = vec![];
        for (rstart, rend) in exon_starts.zip(exon_ends) {
            let start = rstart?;
            let end = rend?;
            res.push((start, end));
        }
        Ok(res)
    }
}

impl From<RefFlatRow> for RefFlatRecord {

    fn from(row: RefFlatRow) -> Self {
        RefFlatRecord {
            gene_id: row.0,
            transcript_name: row.1,
            seq_name: row.2,
            strand: row.3,
            trx_start: row.4,
            trx_end: row.5,
            coding_start: row.6,
            coding_end: row.7,
            num_exons: row.8,
            exon_starts: row.9,
            exon_ends: row.10,
        }
    }
}

pub struct Reader<R: io::Read> {
    inner: csv::Reader<R>,
}

impl<R: io::Read> Reader<R> {

    pub fn from_reader(in_reader: R) -> Reader<R> {
        Reader {
            inner: csv::Reader::from_reader(in_reader)
                .delimiter(b'\t')
                .has_headers(false)
        }
    }

    pub fn records<'a>(&'a mut self) -> RefFlatRecords<'a, R> {
        RefFlatRecords {
            inner: self.inner.decode()
        }
    }

    pub fn transcripts<'a>(&'a mut self) -> RefFlatTranscripts<'a, R> {
        RefFlatTranscripts {
            inner: self.records()
        }
    }

    pub fn genes<'a>(&'a mut self) -> RefFlatGenes<'a, R> {
        RefFlatGenes {
            inner: self.records().group_by(RefFlatGenes::<'a, R>::group_func),
        }
    }
}

impl Reader<fs::File> {
    pub fn from_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::open(path).map(Reader::from_reader)
    }
}

pub struct RefFlatRecords<'a, R: 'a> where R: io::Read {
    inner: csv::DecodedRecords<'a, R, RefFlatRow>,
}

impl<'a, R> Iterator for RefFlatRecords<'a, R> where R: io::Read {

    type Item = Result<RefFlatRecord, Error>;

    fn next(&mut self) -> Option<Result<RefFlatRecord, Error>> {
        self.inner.next()
            .map(|row| {
                row.or_else(|err| Err(Error::from(err)))
                    .map(RefFlatRecord::from)
            })
    }
}

pub struct RefFlatTranscripts<'a, R: 'a> where R: io::Read {
    inner: RefFlatRecords<'a, R>,
}

impl<'a, R> Iterator for RefFlatTranscripts<'a, R> where R: io::Read {

    type Item = Result<Transcript, Error>;

    fn next(&mut self) -> Option<Result<Transcript, Error>> {
        self.inner.next()
            .map(|record| record.and_then(|rec| rec.into_transcript()))
    }
}

type GroupKey = Option<(String, String, char)>;

type GroupFunc = fn(&Result<RefFlatRecord, Error>) -> GroupKey;

type GroupedRecords<'a, 'b, R> = Group<'b, GroupKey, RefFlatRecords<'a, R>, GroupFunc>;

pub struct RefFlatGenes<'a, R: 'a> where R: io::Read, {
    inner: GroupBy<GroupKey, RefFlatRecords<'a, R>, GroupFunc>,
}

impl<'a, R> RefFlatGenes<'a, R> where R: io::Read {

    fn group_func(result: &Result<RefFlatRecord, Error>) -> GroupKey {
        result.as_ref().ok()
            .map(|ref res| (res.gene_id.clone(), res.seq_name.clone(), res.strand.clone()))
    }

    fn group_to_gene<'b>(group: (GroupKey, GroupedRecords<'a, 'b, R>)) -> Result<Gene, Error> {
        let (group_key, records) = group;
        match group_key {

            None => Err(records.filter_map(|x| x.err()).next().unwrap()),

            Some((gid, seq_name, strand_char)) => {
                let mut transcripts = LinkedHashMap::new();
                let (mut gene_start, mut gene_end) = (u64::max_value(), u64::min_value());
                for record in records {
                    let transcript = record.and_then(|rec| rec.into_transcript())?;
                    gene_start = min(gene_start, transcript.start());
                    gene_end = max(gene_end, transcript.end());
                    let tid = transcript.id().map(|id| id.to_owned())
                        .ok_or(Error::RefFlat("transcript does not have ID"))?;
                    transcripts.insert(tid, transcript);
                }
                GBuilder::new(seq_name, gene_start, gene_end)
                    .id(gid)
                    .strand_char(strand_char)
                    .transcripts(transcripts)
                    .transcript_coding_incl_stop(true)
                    .build()
            },
        }
    }
}

impl<'a, R> Iterator for RefFlatGenes<'a, R> where R: io::Read {

    type Item = Result<Gene, Error>;

    fn next(&mut self) -> Option<Result<Gene, Error>> {
        self.inner.into_iter().map(Self::group_to_gene).next()
    }
}

pub struct Writer<W: io::Write> {
    inner: csv::Writer<W>,
}

impl<W: io::Write> Writer<W> {

    pub fn from_writer(in_writer: W) -> Writer<W> {
        Writer {
            inner: csv::Writer::from_writer(in_writer)
                .delimiter(b'\t')
                .quote_style(csv::QuoteStyle::Never)
        }
    }

    pub fn write(&mut self, row: &RefFlatRow) -> Result<(), Error> {
        self.inner
            .encode((&row.0, &row.1, &row.2, row.3, row.4, row.5, row.6, row.7, row.8,
                     &row.9, &row.10))
            .map_err(Error::from)
    }

    pub fn write_record(&mut self, record: &RefFlatRecord) -> Result<(), Error> {
        self.inner
            .encode((&record.gene_id, &record.transcript_name, &record.seq_name,
                     record.strand, record.trx_start, record.trx_end,
                     record.coding_start, record.coding_end, record.num_exons,
                     &record.exon_starts, &record.exon_ends))
            .map_err(Error::from)
    }

    pub fn write_transcript(&mut self, transcript: &Transcript) -> Result<(), Error> {
        let transcript_name = transcript.id()
            .map(|tn| tn.as_ref()).unwrap_or("");
        let strand_char = match transcript.strand() {
            &Strand::Forward => '+',
            &Strand::Reverse => '-',
            &Strand::Unknown => '.',
        };

        let (coding_start, coding_end) = transcript.coding_coord(true)
            .unwrap_or((transcript.end(), transcript.end()));
        let (exon_starts, exon_ends) = Self::coords_field(&transcript);

        self.inner
            .encode((transcript.gene_id(), transcript_name, transcript.seq_name(), strand_char,
                     transcript.start(), transcript.end(),
                     coding_start, coding_end, transcript.exons().len(),
                     exon_starts, exon_ends))
            .map_err(Error::from)
    }

    pub fn write_gene(&mut self, gene: &Gene) -> Result<(), Error> {
        for transcript in gene.transcripts().values() {
            self.write_transcript(&transcript)?;
        }
        Ok(())
    }

    #[inline(always)]
    fn coords_field(trx: &Transcript) -> (String, String) {
        let mut starts = trx.exons().iter().map(|exon| exon.start()).join(",");
        starts.push(',');
        let mut ends = trx.exons().iter().map(|exon| exon.end()).join(",");
        ends.push(',');
        (starts, ends)
    }
}

impl Writer<fs::File> {

    pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
        let f = fs::File::create(path).map_err(Error::from)?;
        Ok(Writer::from_writer(f))
    }
}

impl Writer<Vec<u8>> {

    pub fn from_memory() -> Writer<Vec<u8>> {
        Writer::from_writer(Vec::with_capacity(1024 * 64))
    }

    pub fn as_string(&mut self) -> &str {
        self.inner.as_string()
    }
}
<commit_msg>Use INIT_COORD in io_refflat<commit_after>//! Reader and writer for the refFlat format.
//!
//! The refFlat format is a transcript-oriented format in which each
//! transcript is denoted in a single line. It is used most prominently
//! by the [picard suite tools](https://broadinstitute.github.io/picard/)
//!
//! A minimum specification of the columns can be found here:
//! https://genome.ucsc.edu/goldenPath/gbdDescriptionsOld.html#RefFlat

use std::cmp::{max, min};
use std::convert::AsRef;
use std::io;
use std::fs;
use std::path::Path;
use std::str::FromStr;

use csv;
use itertools::{GroupBy, Group, Itertools};
use linked_hash_map::LinkedHashMap;

use feature::FeatureError;
use {Coord, Gene, GBuilder, Strand, Transcript, TBuilder, Error, INIT_COORD};

pub type RefFlatRow = (String, String, String, char, u64, u64, u64, u64, usize, String, String);

pub struct RefFlatRecord {
    pub gene_id: String,
    pub transcript_name: String,
    pub seq_name: String,
    pub strand: char,
    pub trx_start: u64,
    pub trx_end: u64,
    pub coding_start: u64,
    pub coding_end: u64,
    pub num_exons: usize,
    pub exon_starts: String,
    pub exon_ends: String,
}

impl RefFlatRecord {

    pub fn into_transcript(self) -> Result<Transcript, Error> {

        let exon_coords = self.zip_raw_exon_coords()?;
        if exon_coords.len() != self.num_exons {
            return Err(Error::RefFlat(
                "number of exon and exon coordinates mismatch"));
        }

        let strand = Strand::from_char(&self.strand)
            .map_err(FeatureError::from)
            .map_err(Error::from)?;

        let coding_coord =
            if self.coding_start != self.coding_end {
                Some((self.coding_start, self.coding_end))
            } else {
                None
            };

        TBuilder::new(self.seq_name, self.trx_start, self.trx_end)
            .id(self.transcript_name)
            .gene_id(self.gene_id)
            .strand(strand)
            .coords(exon_coords, coding_coord)
            .coding_incl_stop(true)
            .build()
            .map_err(Error::from)
    }

    #[inline]
    fn zip_raw_exon_coords(&self) -> Result<Vec<Coord<u64>>, Error> {
        let exon_starts = self.exon_starts
            .trim_matches(',').split(',').map(|item| u64::from_str(item).map_err(Error::from));
        let exon_ends = self.exon_ends
            .trim_matches(',').split(',').map(|item| u64::from_str(item).map_err(Error::from));

        let mut res = vec![];
        for (rstart, rend) in exon_starts.zip(exon_ends) {
            let start = rstart?;
            let end = rend?;
            res.push((start, end));
        }
        Ok(res)
    }
}

impl From<RefFlatRow> for RefFlatRecord {

    fn from(row: RefFlatRow) -> Self {
        RefFlatRecord {
            gene_id: row.0,
            transcript_name: row.1,
            seq_name: row.2,
            strand: row.3,
            trx_start: row.4,
            trx_end: row.5,
            coding_start: row.6,
            coding_end: row.7,
            num_exons: row.8,
            exon_starts: row.9,
            exon_ends: row.10,
        }
    }
}

pub struct Reader<R: io::Read> {
    inner: csv::Reader<R>,
}

impl<R: io::Read> Reader<R> {

    pub fn from_reader(in_reader: R) -> Reader<R> {
        Reader {
            inner: csv::Reader::from_reader(in_reader)
                .delimiter(b'\t')
                .has_headers(false)
        }
    }

    pub fn records<'a>(&'a mut self) -> RefFlatRecords<'a, R> {
        RefFlatRecords {
            inner: self.inner.decode()
        }
    }

    pub fn transcripts<'a>(&'a mut self) -> RefFlatTranscripts<'a, R> {
        RefFlatTranscripts {
            inner: self.records()
        }
    }

    pub fn genes<'a>(&'a mut self) -> RefFlatGenes<'a, R> {
        RefFlatGenes {
            inner: self.records().group_by(RefFlatGenes::<'a, R>::group_func),
        }
    }
}

impl Reader<fs::File> {
    pub fn from_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::open(path).map(Reader::from_reader)
    }
}

pub struct RefFlatRecords<'a, R: 'a> where R: io::Read {
    inner: csv::DecodedRecords<'a, R, RefFlatRow>,
}

impl<'a, R> Iterator for RefFlatRecords<'a, R> where R: io::Read {

    type Item = Result<RefFlatRecord, Error>;

    fn next(&mut self) -> Option<Result<RefFlatRecord, Error>> {
        self.inner.next()
            .map(|row| {
                row.or_else(|err| Err(Error::from(err)))
                    .map(RefFlatRecord::from)
            })
    }
}

pub struct RefFlatTranscripts<'a, R: 'a> where R: io::Read {
    inner: RefFlatRecords<'a, R>,
}

impl<'a, R> Iterator for RefFlatTranscripts<'a, R> where R: io::Read {

    type Item = Result<Transcript, Error>;

    fn next(&mut self) -> Option<Result<Transcript, Error>> {
        self.inner.next()
            .map(|record| record.and_then(|rec| rec.into_transcript()))
    }
}

type GroupKey = Option<(String, String, char)>;

type GroupFunc = fn(&Result<RefFlatRecord, Error>) -> GroupKey;

type GroupedRecords<'a, 'b, R> = Group<'b, GroupKey, RefFlatRecords<'a, R>, GroupFunc>;

pub struct RefFlatGenes<'a, R: 'a> where R: io::Read, {
    inner: GroupBy<GroupKey, RefFlatRecords<'a, R>, GroupFunc>,
}

impl<'a, R> RefFlatGenes<'a, R> where R: io::Read {

    fn group_func(result: &Result<RefFlatRecord, Error>) -> GroupKey {
        result.as_ref().ok()
            .map(|ref res| (res.gene_id.clone(), res.seq_name.clone(), res.strand.clone()))
    }

    fn group_to_gene<'b>(group: (GroupKey, GroupedRecords<'a, 'b, R>)) -> Result<Gene, Error> {
        let (group_key, records) = group;
        match group_key {

            None => Err(records.filter_map(|x| x.err()).next().unwrap()),

            Some((gid, seq_name, strand_char)) => {
                let mut transcripts = LinkedHashMap::new();
                let (mut gene_start, mut gene_end) = INIT_COORD;
                for record in records {
                    let transcript = record.and_then(|rec| rec.into_transcript())?;
                    gene_start = min(gene_start, transcript.start());
                    gene_end = max(gene_end, transcript.end());
                    let tid = transcript.id().map(|id| id.to_owned())
                        .ok_or(Error::RefFlat("transcript does not have ID"))?;
                    transcripts.insert(tid, transcript);
                }
                GBuilder::new(seq_name, gene_start, gene_end)
                    .id(gid)
                    .strand_char(strand_char)
                    .transcripts(transcripts)
                    .transcript_coding_incl_stop(true)
                    .build()
            },
        }
    }
}

impl<'a, R> Iterator for RefFlatGenes<'a, R> where R: io::Read {

    type Item = Result<Gene, Error>;

    fn next(&mut self) -> Option<Result<Gene, Error>> {
        self.inner.into_iter().map(Self::group_to_gene).next()
    }
}

pub struct Writer<W: io::Write> {
    inner: csv::Writer<W>,
}

impl<W: io::Write> Writer<W> {

    pub fn from_writer(in_writer: W) -> Writer<W> {
        Writer {
            inner: csv::Writer::from_writer(in_writer)
                .delimiter(b'\t')
                .quote_style(csv::QuoteStyle::Never)
        }
    }

    pub fn write(&mut self, row: &RefFlatRow) -> Result<(), Error> {
        self.inner
            .encode((&row.0, &row.1, &row.2, row.3, row.4, row.5, row.6, row.7, row.8,
                     &row.9, &row.10))
            .map_err(Error::from)
    }

    pub fn write_record(&mut self, record: &RefFlatRecord) -> Result<(), Error> {
        self.inner
            .encode((&record.gene_id, &record.transcript_name, &record.seq_name,
                     record.strand, record.trx_start, record.trx_end,
                     record.coding_start, record.coding_end, record.num_exons,
                     &record.exon_starts, &record.exon_ends))
            .map_err(Error::from)
    }

    pub fn write_transcript(&mut self, transcript: &Transcript) -> Result<(), Error> {
        let transcript_name = transcript.id()
            .map(|tn| tn.as_ref()).unwrap_or("");
        let strand_char = match transcript.strand() {
            &Strand::Forward => '+',
            &Strand::Reverse => '-',
            &Strand::Unknown => '.',
        };

        let (coding_start, coding_end) = transcript.coding_coord(true)
            .unwrap_or((transcript.end(), transcript.end()));
        let (exon_starts, exon_ends) = Self::coords_field(&transcript);

        self.inner
            .encode((transcript.gene_id(), transcript_name, transcript.seq_name(), strand_char,
                     transcript.start(), transcript.end(),
                     coding_start, coding_end, transcript.exons().len(),
                     exon_starts, exon_ends))
            .map_err(Error::from)
    }

    pub fn write_gene(&mut self, gene: &Gene) -> Result<(), Error> {
        for transcript in gene.transcripts().values() {
            self.write_transcript(&transcript)?;
        }
        Ok(())
    }

    #[inline(always)]
    fn coords_field(trx: &Transcript) -> (String, String) {
        let mut starts = trx.exons().iter().map(|exon| exon.start()).join(",");
        starts.push(',');
        let mut ends = trx.exons().iter().map(|exon| exon.end()).join(",");
        ends.push(',');
        (starts, ends)
    }
}

impl Writer<fs::File> {

    pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
        let f = fs::File::create(path).map_err(Error::from)?;
        Ok(Writer::from_writer(f))
    }
}

impl Writer<Vec<u8>> {

    pub fn from_memory() -> Writer<Vec<u8>> {
        Writer::from_writer(Vec::with_capacity(1024 * 64))
    }

    pub fn as_string(&mut self) -> &str {
        self.inner.as_string()
    }
}
<|endoftext|>


<commit_before>use crate::codec::*;
use crate::endpoint::Endpoint;
use crate::error::ZmqResult;
use crate::message::*;
use crate::transport::AcceptStopHandle;
use crate::util::PeerIdentity;
use crate::{
    MultiPeerBackend, Socket, SocketBackend, SocketEvent, SocketOptions, SocketRecv, SocketType,
};

use crate::backend::Peer;
use crate::fair_queue::FairQueue;
use crate::fair_queue::QueueInner;
use async_trait::async_trait;
use bytes::{BufMut, BytesMut};
use crossbeam::queue::SegQueue;
use dashmap::DashMap;
use futures::channel::mpsc;
use futures::{SinkExt, StreamExt};
use parking_lot::Mutex;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;

pub enum SubBackendMsgType {
    UNSUBSCRIBE = 0,
    SUBSCRIBE = 1,
}

pub(crate) struct SubSocketBackend {
    pub(crate) peers: DashMap<PeerIdentity, Peer>,
    fair_queue_inner: Option<Arc<Mutex<QueueInner<ZmqFramedRead, PeerIdentity>>>>,
    pub(crate) round_robin: SegQueue<PeerIdentity>,
    socket_type: SocketType,
    socket_options: SocketOptions,
    pub(crate) socket_monitor: Mutex<Option<mpsc::Sender<SocketEvent>>>,
    subs: Mutex<HashSet<String>>,
}

impl SubSocketBackend {
    pub(crate) fn with_options(
        fair_queue_inner: Option<Arc<Mutex<QueueInner<ZmqFramedRead, PeerIdentity>>>>,
        socket_type: SocketType,
        options: SocketOptions,
    ) -> Self {
        Self {
            peers: DashMap::new(),
            fair_queue_inner,
            round_robin: SegQueue::new(),
            socket_type,
            socket_options: options,
            socket_monitor: Mutex::new(None),
            subs: Mutex::new(HashSet::new()),
        }
    }

    pub fn create_subs_message(subscription: &str, msg_type: SubBackendMsgType) -> ZmqMessage {
        let mut buf = BytesMut::with_capacity(subscription.len() + 1);
        buf.put_u8(msg_type as u8);
        buf.extend_from_slice(subscription.as_bytes());

        buf.freeze().into()
    }
}

impl SocketBackend for SubSocketBackend {
    fn socket_type(&self) -> SocketType {
        self.socket_type
    }

    fn socket_options(&self) -> &SocketOptions {
        &self.socket_options
    }

    fn shutdown(&self) {
        self.peers.clear();
    }

    fn monitor(&self) -> &Mutex<Option<mpsc::Sender<SocketEvent>>> {
        &self.socket_monitor
    }
}

#[async_trait]
impl MultiPeerBackend for SubSocketBackend {
    async fn peer_connected(self: Arc<Self>, peer_id: &PeerIdentity, io: FramedIo) {
        let (recv_queue, mut send_queue) = io.into_parts();

        let subs_msgs: Vec<ZmqMessage> = self
            .subs
            .lock()
            .iter()
            .map(|x| SubSocketBackend::create_subs_message(x, SubBackendMsgType::SUBSCRIBE))
            .collect();

        for message in subs_msgs {
            send_queue
                .send(Message::Message(message))
                .await
                .unwrap();
        }

        self.peers.insert(peer_id.clone(), Peer { send_queue });
        self.round_robin.push(peer_id.clone());
        match &self.fair_queue_inner {
            None => {}
            Some(inner) => {
                inner.lock().insert(peer_id.clone(), recv_queue);
            }
        };
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        self.peers.remove(peer_id);
    }
}

pub struct SubSocket {
    backend: Arc<SubSocketBackend>,
    fair_queue: FairQueue<ZmqFramedRead, PeerIdentity>,
    binds: HashMap<Endpoint, AcceptStopHandle>,
}

impl Drop for SubSocket {
    fn drop(&mut self) {
        self.backend.shutdown()
    }
}

impl SubSocket {
    pub async fn subscribe(&mut self, subscription: &str) -> ZmqResult<()> {
        self.backend.subs.lock().insert(subscription.to_string());
        self.process_subs(subscription, SubBackendMsgType::SUBSCRIBE)
            .await
    }

    pub async fn unsubscribe(&mut self, subscription: &str) -> ZmqResult<()> {
        self.backend.subs.lock().remove(subscription);
        self.process_subs(subscription, SubBackendMsgType::UNSUBSCRIBE)
            .await
    }

    async fn process_subs(
        &mut self,
        subscription: &str,
        msg_type: SubBackendMsgType,
    ) -> ZmqResult<()> {
        let message: ZmqMessage = SubSocketBackend::create_subs_message(subscription, msg_type);

        for mut peer in self.backend.peers.iter_mut() {
            peer.send_queue
                .send(Message::Message(message.clone()))
                .await?;
        }
        Ok(())
    }
}

#[async_trait]
impl Socket for SubSocket {
    fn with_options(options: SocketOptions) -> Self {
        let fair_queue = FairQueue::new(true);
        Self {
            backend: Arc::new(SubSocketBackend::with_options(
                Some(fair_queue.inner()),
                SocketType::SUB,
                options,
            )),
            fair_queue,
            binds: HashMap::new(),
        }
    }

    fn backend(&self) -> Arc<dyn MultiPeerBackend> {
        self.backend.clone()
    }

    fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle> {
        &mut self.binds
    }

    fn monitor(&mut self) -> mpsc::Receiver<SocketEvent> {
        let (sender, receiver) = mpsc::channel(1024);
        self.backend.socket_monitor.lock().replace(sender);
        receiver
    }
}

#[async_trait]
impl SocketRecv for SubSocket {
    async fn recv(&mut self) -> ZmqResult<ZmqMessage> {
        loop {
            match self.fair_queue.next().await {
                Some((_peer_id, Ok(Message::Message(message)))) => {
                    return Ok(message);
                }
                Some((_peer_id, Ok(msg))) => todo!("Unimplemented message: {:?}", msg),
                Some((peer_id, Err(_))) => {
                    self.backend.peer_disconnected(&peer_id);
                }
                None => todo!(),
            }
        }
    }
}
<commit_msg>formatting<commit_after>use crate::codec::*;
use crate::endpoint::Endpoint;
use crate::error::ZmqResult;
use crate::message::*;
use crate::transport::AcceptStopHandle;
use crate::util::PeerIdentity;
use crate::{
    MultiPeerBackend, Socket, SocketBackend, SocketEvent, SocketOptions, SocketRecv, SocketType,
};

use crate::backend::Peer;
use crate::fair_queue::FairQueue;
use crate::fair_queue::QueueInner;
use async_trait::async_trait;
use bytes::{BufMut, BytesMut};
use crossbeam::queue::SegQueue;
use dashmap::DashMap;
use futures::channel::mpsc;
use futures::{SinkExt, StreamExt};
use parking_lot::Mutex;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;

pub enum SubBackendMsgType {
    UNSUBSCRIBE = 0,
    SUBSCRIBE = 1,
}

pub(crate) struct SubSocketBackend {
    pub(crate) peers: DashMap<PeerIdentity, Peer>,
    fair_queue_inner: Option<Arc<Mutex<QueueInner<ZmqFramedRead, PeerIdentity>>>>,
    pub(crate) round_robin: SegQueue<PeerIdentity>,
    socket_type: SocketType,
    socket_options: SocketOptions,
    pub(crate) socket_monitor: Mutex<Option<mpsc::Sender<SocketEvent>>>,
    subs: Mutex<HashSet<String>>,
}

impl SubSocketBackend {
    pub(crate) fn with_options(
        fair_queue_inner: Option<Arc<Mutex<QueueInner<ZmqFramedRead, PeerIdentity>>>>,
        socket_type: SocketType,
        options: SocketOptions,
    ) -> Self {
        Self {
            peers: DashMap::new(),
            fair_queue_inner,
            round_robin: SegQueue::new(),
            socket_type,
            socket_options: options,
            socket_monitor: Mutex::new(None),
            subs: Mutex::new(HashSet::new()),
        }
    }

    pub fn create_subs_message(subscription: &str, msg_type: SubBackendMsgType) -> ZmqMessage {
        let mut buf = BytesMut::with_capacity(subscription.len() + 1);
        buf.put_u8(msg_type as u8);
        buf.extend_from_slice(subscription.as_bytes());

        buf.freeze().into()
    }
}

impl SocketBackend for SubSocketBackend {
    fn socket_type(&self) -> SocketType {
        self.socket_type
    }

    fn socket_options(&self) -> &SocketOptions {
        &self.socket_options
    }

    fn shutdown(&self) {
        self.peers.clear();
    }

    fn monitor(&self) -> &Mutex<Option<mpsc::Sender<SocketEvent>>> {
        &self.socket_monitor
    }
}

#[async_trait]
impl MultiPeerBackend for SubSocketBackend {
    async fn peer_connected(self: Arc<Self>, peer_id: &PeerIdentity, io: FramedIo) {
        let (recv_queue, mut send_queue) = io.into_parts();

        let subs_msgs: Vec<ZmqMessage> = self
            .subs
            .lock()
            .iter()
            .map(|x| SubSocketBackend::create_subs_message(x, SubBackendMsgType::SUBSCRIBE))
            .collect();

        for message in subs_msgs {
            send_queue.send(Message::Message(message)).await.unwrap();
        }

        self.peers.insert(peer_id.clone(), Peer { send_queue });
        self.round_robin.push(peer_id.clone());
        match &self.fair_queue_inner {
            None => {}
            Some(inner) => {
                inner.lock().insert(peer_id.clone(), recv_queue);
            }
        };
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        self.peers.remove(peer_id);
    }
}

pub struct SubSocket {
    backend: Arc<SubSocketBackend>,
    fair_queue: FairQueue<ZmqFramedRead, PeerIdentity>,
    binds: HashMap<Endpoint, AcceptStopHandle>,
}

impl Drop for SubSocket {
    fn drop(&mut self) {
        self.backend.shutdown()
    }
}

impl SubSocket {
    pub async fn subscribe(&mut self, subscription: &str) -> ZmqResult<()> {
        self.backend.subs.lock().insert(subscription.to_string());
        self.process_subs(subscription, SubBackendMsgType::SUBSCRIBE)
            .await
    }

    pub async fn unsubscribe(&mut self, subscription: &str) -> ZmqResult<()> {
        self.backend.subs.lock().remove(subscription);
        self.process_subs(subscription, SubBackendMsgType::UNSUBSCRIBE)
            .await
    }

    async fn process_subs(
        &mut self,
        subscription: &str,
        msg_type: SubBackendMsgType,
    ) -> ZmqResult<()> {
        let message: ZmqMessage = SubSocketBackend::create_subs_message(subscription, msg_type);

        for mut peer in self.backend.peers.iter_mut() {
            peer.send_queue
                .send(Message::Message(message.clone()))
                .await?;
        }
        Ok(())
    }
}

#[async_trait]
impl Socket for SubSocket {
    fn with_options(options: SocketOptions) -> Self {
        let fair_queue = FairQueue::new(true);
        Self {
            backend: Arc::new(SubSocketBackend::with_options(
                Some(fair_queue.inner()),
                SocketType::SUB,
                options,
            )),
            fair_queue,
            binds: HashMap::new(),
        }
    }

    fn backend(&self) -> Arc<dyn MultiPeerBackend> {
        self.backend.clone()
    }

    fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle> {
        &mut self.binds
    }

    fn monitor(&mut self) -> mpsc::Receiver<SocketEvent> {
        let (sender, receiver) = mpsc::channel(1024);
        self.backend.socket_monitor.lock().replace(sender);
        receiver
    }
}

#[async_trait]
impl SocketRecv for SubSocket {
    async fn recv(&mut self) -> ZmqResult<ZmqMessage> {
        loop {
            match self.fair_queue.next().await {
                Some((_peer_id, Ok(Message::Message(message)))) => {
                    return Ok(message);
                }
                Some((_peer_id, Ok(msg))) => todo!("Unimplemented message: {:?}", msg),
                Some((peer_id, Err(_))) => {
                    self.backend.peer_disconnected(&peer_id);
                }
                None => todo!(),
            }
        }
    }
}
<|endoftext|>


<commit_before>//! This module handles the "Universal Chess Interface" protocol
//! communication.

use regex::Regex;
use std::io;
use std::io::{Read, Write, BufRead, BufReader, BufWriter, ErrorKind};

/// A command from the GUI to the engine.
pub enum UciCommand {
    Uci,
    Debug(bool),
    IsReady,
    SetOption(SetOptionParams),
    UciNewGame,
    Position(PositionParams),
    Go(GoParams),
    Stop,
    PonderHit,
    Quit,
}

pub struct SetOptionParams {
    name: String,
    value: String,
}

pub struct PositionParams {
    fen: String,
    moves: Vec<String>,
}

pub struct GoParams {
    searchmoves: Option<Vec<String>>,
    ponder: bool,
    wtime: Option<u64>,
    btime: Option<u64>,
    winc: Option<u64>,
    binc: Option<u64>,
    movestogo: Option<u64>,
    depth: Option<u64>,
    nodes: Option<u64>,
    mate: Option<u64>,
    movetime: Option<u64>,
    infinite: bool,
}

fn parse_uci_command(s: &str) -> Option<UciCommand> {
    const COMMAND: &'static str = "uci|debug|isready|setoption|register|\
                                   ucinewgame|position|go|stop|ponderhit|quit";
    lazy_static! {
        static ref RE: Regex = Regex::new(
            format!(r"\b({})\s*(?:\s(.*)|$)", COMMAND).as_str()
        ).unwrap();
    }

    if let Some(captures) = RE.captures(s) {
        let command = captures.at(1).unwrap();
        let the_rest = captures.at(2).unwrap_or("");
        match command {
            "uci" if the_rest == "" => Some(UciCommand::Uci),
            "stop" if the_rest == "" => Some(UciCommand::Stop),
            "quit" if the_rest == "" => Some(UciCommand::Quit),
            "isready" if the_rest == "" => Some(UciCommand::IsReady),
            "ponderhit" if the_rest == "" => Some(UciCommand::PonderHit),
            "ucinewgame" if the_rest == "" => Some(UciCommand::UciNewGame),
            "debug" => {
                if let Some(params) = parse_debug_params(the_rest) {
                    Some(UciCommand::Debug(params))
                } else {
                    None
                }
            }
            "setoption" => {
                if let Some(params) = parse_setoption_params(the_rest) {
                    Some(UciCommand::SetOption(params))
                } else {
                    None
                }
            }
            "postition" => {
                if let Some(params) = parse_position_params(the_rest) {
                    Some(UciCommand::Position(params))
                } else {
                    None
                }
            }
            "go" => Some(UciCommand::Go(parse_go_params(the_rest))),
            _ => None,
        }
    } else {
        None
    }
}

fn parse_debug_params(s: &str) -> Option<bool> {
    lazy_static! {
        static ref RE: Regex = Regex::new(
            r"^(on|off)\s*$").unwrap();
    }

    if let Some(captures) = RE.captures(s) {
        match captures.at(1).unwrap() {
            "on" => Some(true),
            "off" => Some(false),
            _ => None,
        }
    } else {
        None
    }
}

fn parse_setoption_params(s: &str) -> Option<SetOptionParams> {
    lazy_static! {
        static ref RE: Regex = Regex::new(
            r"^name\s+(.*?)(?:\s+value\s+(.*?))?\s*$").unwrap();
    }

    if let Some(captures) = RE.captures(s) {
        Some(SetOptionParams {
            name: captures.at(1).unwrap().to_string(),
            value: captures.at(2).unwrap_or("").to_string(),
        })
    } else {
        None
    }
}

fn parse_position_params(s: &str) -> Option<PositionParams> {
    const MOVES: &'static str = r"(?:\s+[a-h][1-8][a-h][1-8][qrbn]?)*";
    const STARTPOS: &'static str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 1";
    lazy_static! {
        static ref RE: Regex = Regex::new(
            format!(
                r"^(?:fen\s+(?P<fen>.*?)|startpos)(?:\s+moves(?P<moves>{}))?\s*$",
                MOVES,
            ).as_str()
        ).unwrap();
    }

    if let Some(captures) = RE.captures(s) {
        Some(PositionParams {
            fen: if let Some(fen) = captures.name("fen") {
                fen.to_string()
            } else {
                STARTPOS.to_string()
            },
            moves: captures.name("moves")
                           .unwrap_or("")
                           .split_whitespace()
                           .map(|x| x.to_string())
                           .collect(),
        })
    } else {
        None
    }
}

fn parse_go_params(s: &str) -> GoParams {
    const MOVES: &'static str = r"(?:\s+[a-h][1-8][a-h][1-8][qrbn]?)+";
    const KEYWORD: &'static str = "wtime|btime|winc|binc|movestogo|depth|nodes|\
                                   mate|movetime|ponder|infinite|searchmoves";
    lazy_static! {
        static ref RE: Regex = Regex::new(
            format!(
                r"\b(?P<keyword>{})(?:\s+(?P<number>\d+)|(?P<moves>{}))?(?:\s+|$)",
                KEYWORD,
                MOVES,
            ).as_str()
        ).unwrap();
    }

    let mut params = GoParams {
        searchmoves: None,
        ponder: false,
        wtime: None,
        btime: None,
        winc: None,
        binc: None,
        movestogo: None,
        depth: None,
        nodes: None,
        mate: None,
        movetime: None,
        infinite: false,
    };

    for captures in RE.captures_iter(s) {
        let keyword = captures.name("keyword").unwrap();
        match keyword {
            "searchmoves" => {
                if let Some(moves) = captures.name("moves") {
                    params.searchmoves = Some(moves.split_whitespace()
                                                   .map(|x| x.to_string())
                                                   .collect());
                }
            }
            "infinite" => {
                params.infinite = true;
            }
            "ponder" => {
                params.ponder = true;
            }
            _ => {
                if let Some(number) = captures.name("number") {
                    let field = match keyword {
                        "wtime" => &mut params.wtime,
                        "btime" => &mut params.btime,
                        "winc" => &mut params.winc,
                        "binc" => &mut params.binc,
                        "movestogo" => &mut params.movestogo,
                        "depth" => &mut params.depth,
                        "nodes" => &mut params.nodes,
                        "mate" => &mut params.mate,
                        "movetime" => &mut params.movetime,
                        _ => panic!("invalid keyword"),
                    };
                    *field = number.parse::<u64>().ok();
                }
            }
        }
    }
    params
}

/// A response from the engine to the GUI .
pub enum UciResponse {
    Id {
        attribute: String,
        value: String,
    },
    UciOk,
    ReadyOk,
    BestMove {
        best_move: String,
        ponder: Option<String>,
    },
    Info(String),
    Option(OptionDescription),
}

/// A description of a single configuration option (name and value)
/// supported by the engine.
///
/// The GUI may use this information to configure the engine. It may
/// also build dialog boxes according to the received option
/// descriptions so that GUI users can configure the engine.
pub struct OptionDescription {
    name: String,
    description: ValueDescription,
}

/// A description of a single configurable value.
pub enum ValueDescription {
    Check {
        default: bool,
    },
    Spin {
        min: i32,
        max: i32,
        default: i32,
    },
    Combo {
        list: Vec<String>,
        default: String,
    },
    String {
        default: String,
    },
    Button,
}

/// The main UCI protocol serving loop.
pub struct UciServingLoop<R: Read, W: Write> {
    reader: BufReader<R>,
    writer: BufWriter<W>,
    // engine: Box<UciEngine>,
    engine_is_started: bool,
    engine_is_thinking: bool,
}

impl<R: Read, W: Write> UciServingLoop<R, W> {
    fn wait_for_hanshake(in_stream: R, out_stream: W) -> io::Result<Self> {
        let mut reader = BufReader::new(in_stream);
        let mut writer = BufWriter::new(out_stream);
        let mut line = String::new();
        if try!(reader.read_line(&mut line)) == 0 {
            return Err(io::Error::new(ErrorKind::UnexpectedEof, "EOF"));
        }
        if line.as_str() != "uci" {
            return Err(io::Error::new(ErrorKind::Other, "unrecognized protocol"));
        }
        Ok(UciServingLoop {
            reader: reader,
            writer: writer,
            engine_is_started: false,
            engine_is_thinking: false,
        })
    }
}

pub trait UciEngine {
    fn options(&self) -> Vec<OptionDescription>;
}

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

    #[test]
    fn test_parse_go_params() {
        use super::parse_go_params;
        assert_eq!(parse_go_params(" wtime22000  ").wtime, None);
        assert_eq!(parse_go_params(" wtime    22000  ").wtime, Some(22000));
        assert_eq!(parse_go_params("wtime 22000").wtime, Some(22000));
        assert_eq!(parse_go_params("wtime 99999999999999998888888888999999999999999999").wtime,
                   None);
        assert_eq!(parse_go_params("wtime 22000").infinite, false);
        assert_eq!(parse_go_params("searchmoves   e2e4  c7c8q  ").searchmoves,
                   Some(vec!["e2e4".to_string(), "c7c8q".to_string()]));
        assert_eq!(parse_go_params("searchmoves   e2e4  c7c8q,ponder  ").searchmoves,
                   Some(vec!["e2e4".to_string()]));
        assert_eq!(parse_go_params("searchmoves aabb").searchmoves, None);
        assert_eq!(parse_go_params("infinite wtime 22000").wtime, Some(22000));
        assert_eq!(parse_go_params("infinite wtime 22000").infinite, true);
        assert_eq!(parse_go_params("wtime 22000 infinite btime 11000").infinite,
                   true);
        assert_eq!(parse_go_params("wtime fdfee / 22000 infinite btime 11000 fdfds").infinite,
                   true);
        assert_eq!(parse_go_params("wtime 22000 infinite btime 11000 ponder").btime,
                   Some(11000));
    }

    #[test]
    fn test_parse_debug_params() {
        use super::parse_debug_params;
        assert_eq!(parse_debug_params("on"), Some(true));
        assert_eq!(parse_debug_params(" on"), None);
        assert_eq!(parse_debug_params("off   "), Some(false));
        assert_eq!(parse_debug_params(" off   "), None);
        assert_eq!(parse_debug_params("    "), None);
    }

    #[test]
    fn test_parse_setoption_params() {
        use super::parse_setoption_params;
        assert_eq!(parse_setoption_params("name   xxx  value   yyy  ").unwrap().name,
                   "xxx".to_string());
        assert_eq!(parse_setoption_params("name xxx value yyy").unwrap().value,
                   "yyy".to_string());
        assert_eq!(parse_setoption_params("name xxx   value  ").unwrap().value,
                   "".to_string());
        assert_eq!(parse_setoption_params("name xxx    ").unwrap().value,
                   "".to_string());
        assert_eq!(parse_setoption_params("name     ").unwrap().name,
                   "".to_string());
        assert_eq!(parse_setoption_params("name     ").unwrap().value,
                   "".to_string());
        assert!(parse_setoption_params("namexxx     ").is_none());
    }

    #[test]
    fn test_parse_position_params() {
        use super::parse_position_params;
        assert_eq!(parse_position_params("startpos  ").unwrap().fen,
                   "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 1");
        assert_eq!(parse_position_params("startpos ").unwrap().moves.len(), 0);
        assert_eq!(parse_position_params("startpos   moves  ").unwrap().moves.len(),
                   0);
        assert_eq!(parse_position_params("startpos   moves   e2e4   d2d4 ").unwrap().moves.len(),
                   2);
        assert_eq!(parse_position_params("fen xxx moves e2e4").unwrap().moves.len(),
                   1);
        assert_eq!(parse_position_params("fen xxx moves e2e4").unwrap().fen,
                   "xxx".to_string());
        assert_eq!(parse_position_params("fen   xxx    moves e2e4").unwrap().fen,
                   "xxx".to_string());
        assert_eq!(parse_position_params("fen   xxx    moves").unwrap().fen,
                   "xxx".to_string());
        assert_eq!(parse_position_params("fen   xxx   ").unwrap().fen,
                   "xxx".to_string());
    }
}
<commit_msg>refactoring<commit_after>//! This module handles the "Universal Chess Interface" protocol
//! communication.

use regex::Regex;
use std::io;
use std::io::{Read, Write, BufRead, BufReader, BufWriter, ErrorKind};

/// A command from the GUI to the engine.
pub enum UciCommand {
    Uci,
    Debug(bool),
    IsReady,
    SetOption(SetOptionParams),
    UciNewGame,
    Position(PositionParams),
    Go(GoParams),
    Stop,
    PonderHit,
    Quit,
}

pub struct SetOptionParams {
    name: String,
    value: String,
}

pub struct PositionParams {
    fen: String,
    moves: Vec<String>,
}

pub struct GoParams {
    searchmoves: Option<Vec<String>>,
    ponder: bool,
    wtime: Option<u64>,
    btime: Option<u64>,
    winc: Option<u64>,
    binc: Option<u64>,
    movestogo: Option<u64>,
    depth: Option<u64>,
    nodes: Option<u64>,
    mate: Option<u64>,
    movetime: Option<u64>,
    infinite: bool,
}

struct ParseError;

fn parse_uci_command(s: &str) -> Option<UciCommand> {
    lazy_static! {
        static ref RE: Regex = Regex::new(
            format!(r"\b({})\s*(?:\s(.*)|$)",
                    "uci|debug|isready|setoption|register|\
                     ucinewgame|position|go|stop|ponderhit|quit",  // any UCI command
            ).as_str()
        ).unwrap();
    }
    if let Some(captures) = RE.captures(s) {
        let command_str = captures.at(1).unwrap();
        let params_str = captures.at(2).unwrap_or("");
        match command_str {
            "uci" if params_str == "" => Some(UciCommand::Uci),
            "stop" if params_str == "" => Some(UciCommand::Stop),
            "quit" if params_str == "" => Some(UciCommand::Quit),
            "isready" if params_str == "" => Some(UciCommand::IsReady),
            "ponderhit" if params_str == "" => Some(UciCommand::PonderHit),
            "ucinewgame" if params_str == "" => Some(UciCommand::UciNewGame),
            "debug" => {
                if let Some(p) = parse_debug_params(params_str) {
                    Some(UciCommand::Debug(p))
                } else {
                    None
                }
            }
            "setoption" => {
                if let Some(p) = parse_setoption_params(params_str) {
                    Some(UciCommand::SetOption(p))
                } else {
                    None
                }
            }
            "postition" => {
                if let Some(p) = parse_position_params(params_str) {
                    Some(UciCommand::Position(p))
                } else {
                    None
                }
            }
            "go" => Some(UciCommand::Go(parse_go_params(params_str))),
            _ => None,
        }
    } else {
        None
    }
}

fn parse_debug_params(s: &str) -> Option<bool> {
    lazy_static! {
        static ref RE: Regex = Regex::new(
            r"^(on|off)\s*$").unwrap();
    }
    if let Some(captures) = RE.captures(s) {
        match captures.at(1).unwrap() {
            "on" => Some(true),
            "off" => Some(false),
            _ => None,
        }
    } else {
        None
    }
}

fn parse_setoption_params(s: &str) -> Option<SetOptionParams> {
    lazy_static! {
        static ref RE: Regex = Regex::new(
            r"^name\s+(.*?)(?:\s+value\s+(.*?))?\s*$").unwrap();
    }
    if let Some(captures) = RE.captures(s) {
        Some(SetOptionParams {
            name: captures.at(1).unwrap().to_string(),
            value: captures.at(2).unwrap_or("").to_string(),
        })
    } else {
        None
    }
}

fn parse_position_params(s: &str) -> Option<PositionParams> {
    const STARTPOS: &'static str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 1";
    lazy_static! {
        static ref RE: Regex = Regex::new(
            format!(
                r"^(?:fen\s+(?P<fen>.*?)|startpos)(?:\s+moves(?P<moves>{}))?\s*$",
                r"(?:\s+[a-h][1-8][a-h][1-8][qrbn]?)*",  // a possibly empty list of moves
            ).as_str()
        ).unwrap();
    }
    if let Some(captures) = RE.captures(s) {
        Some(PositionParams {
            fen: if let Some(fen) = captures.name("fen") {
                fen.to_string()
            } else {
                STARTPOS.to_string()
            },
            moves: captures.name("moves")
                           .unwrap_or("")
                           .split_whitespace()
                           .map(|x| x.to_string())
                           .collect(),
        })
    } else {
        None
    }
}

fn parse_go_params(s: &str) -> GoParams {
    lazy_static! {
        static ref RE: Regex = Regex::new(
            format!(
                r"\b(?P<keyword>{})(?:\s+(?P<number>\d+)|(?P<moves>{}))?(?:\s+|$)",
                "wtime|btime|winc|binc|movestogo|depth|\
                 nodes|mate|movetime|ponder|infinite|searchmoves",  // any keyword
                r"(?:\s+[a-h][1-8][a-h][1-8][qrbn]?)+",  // a non-empty list of moves
            ).as_str()
        ).unwrap();
    }
    let mut params = GoParams {
        searchmoves: None,
        ponder: false,
        wtime: None,
        btime: None,
        winc: None,
        binc: None,
        movestogo: None,
        depth: None,
        nodes: None,
        mate: None,
        movetime: None,
        infinite: false,
    };
    for captures in RE.captures_iter(s) {
        let keyword = captures.name("keyword").unwrap();
        match keyword {
            "searchmoves" => {
                if let Some(moves) = captures.name("moves") {
                    params.searchmoves = Some(moves.split_whitespace()
                                                   .map(|x| x.to_string())
                                                   .collect());
                }
            }
            "infinite" => {
                params.infinite = true;
            }
            "ponder" => {
                params.ponder = true;
            }
            _ => {
                if let Some(number) = captures.name("number") {
                    let field = match keyword {
                        "wtime" => &mut params.wtime,
                        "btime" => &mut params.btime,
                        "winc" => &mut params.winc,
                        "binc" => &mut params.binc,
                        "movestogo" => &mut params.movestogo,
                        "depth" => &mut params.depth,
                        "nodes" => &mut params.nodes,
                        "mate" => &mut params.mate,
                        "movetime" => &mut params.movetime,
                        _ => panic!("invalid keyword"),
                    };
                    *field = number.parse::<u64>().ok();
                }
            }
        }
    }
    params
}

/// A response from the engine to the GUI .
pub enum UciResponse {
    Id {
        attribute: String,
        value: String,
    },
    UciOk,
    ReadyOk,
    BestMove {
        best_move: String,
        ponder: Option<String>,
    },
    Info(String),
    Option(OptionDescription),
}

/// A description of a single configuration option (name and value)
/// supported by the engine.
///
/// The GUI may use this information to configure the engine. It may
/// also build dialog boxes according to the received option
/// descriptions so that GUI users can configure the engine.
pub struct OptionDescription {
    name: String,
    description: ValueDescription,
}

/// A description of a single configurable value.
pub enum ValueDescription {
    Check {
        default: bool,
    },
    Spin {
        min: i32,
        max: i32,
        default: i32,
    },
    Combo {
        list: Vec<String>,
        default: String,
    },
    String {
        default: String,
    },
    Button,
}

/// The main UCI protocol serving loop.
pub struct UciServingLoop<R: Read, W: Write> {
    reader: BufReader<R>,
    writer: BufWriter<W>,
    // engine: Box<UciEngine>,
    engine_is_started: bool,
    engine_is_thinking: bool,
}

impl<R: Read, W: Write> UciServingLoop<R, W> {
    fn wait_for_hanshake(in_stream: R, out_stream: W) -> io::Result<Self> {
        let mut reader = BufReader::new(in_stream);
        let mut writer = BufWriter::new(out_stream);
        let mut line = String::new();
        if try!(reader.read_line(&mut line)) == 0 {
            return Err(io::Error::new(ErrorKind::UnexpectedEof, "EOF"));
        }
        if line.as_str() != "uci" {
            return Err(io::Error::new(ErrorKind::Other, "unrecognized protocol"));
        }
        Ok(UciServingLoop {
            reader: reader,
            writer: writer,
            engine_is_started: false,
            engine_is_thinking: false,
        })
    }
}

pub trait UciEngine {
    fn options(&self) -> Vec<OptionDescription>;
}

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

    #[test]
    fn test_parse_go_params() {
        use super::parse_go_params;
        assert_eq!(parse_go_params(" wtime22000  ").wtime, None);
        assert_eq!(parse_go_params(" wtime    22000  ").wtime, Some(22000));
        assert_eq!(parse_go_params("wtime 22000").wtime, Some(22000));
        assert_eq!(parse_go_params("wtime 99999999999999998888888888999999999999999999").wtime,
                   None);
        assert_eq!(parse_go_params("wtime 22000").infinite, false);
        assert_eq!(parse_go_params("searchmoves   e2e4  c7c8q  ").searchmoves,
                   Some(vec!["e2e4".to_string(), "c7c8q".to_string()]));
        assert_eq!(parse_go_params("searchmoves   e2e4  c7c8q,ponder  ").searchmoves,
                   Some(vec!["e2e4".to_string()]));
        assert_eq!(parse_go_params("searchmoves aabb").searchmoves, None);
        assert_eq!(parse_go_params("infinite wtime 22000").wtime, Some(22000));
        assert_eq!(parse_go_params("infinite wtime 22000").infinite, true);
        assert_eq!(parse_go_params("wtime 22000 infinite btime 11000").infinite,
                   true);
        assert_eq!(parse_go_params("wtime fdfee / 22000 infinite btime 11000 fdfds").infinite,
                   true);
        assert_eq!(parse_go_params("wtime 22000 infinite btime 11000 ponder").btime,
                   Some(11000));
    }

    #[test]
    fn test_parse_debug_params() {
        use super::parse_debug_params;
        assert_eq!(parse_debug_params("on"), Some(true));
        assert_eq!(parse_debug_params(" on"), None);
        assert_eq!(parse_debug_params("off   "), Some(false));
        assert_eq!(parse_debug_params(" off   "), None);
        assert_eq!(parse_debug_params("    "), None);
        assert_eq!(parse_debug_params(" sgssgs  "), None);
    }

    #[test]
    fn test_parse_setoption_params() {
        use super::parse_setoption_params;
        assert_eq!(parse_setoption_params("name   xxx  value   yyy  ").unwrap().name,
                   "xxx".to_string());
        assert_eq!(parse_setoption_params("name xxx value yyy").unwrap().value,
                   "yyy".to_string());
        assert_eq!(parse_setoption_params("name xxx   value  ").unwrap().value,
                   "".to_string());
        assert_eq!(parse_setoption_params("name xxx    ").unwrap().value,
                   "".to_string());
        assert_eq!(parse_setoption_params("name     ").unwrap().name,
                   "".to_string());
        assert_eq!(parse_setoption_params("name     ").unwrap().value,
                   "".to_string());
        assert!(parse_setoption_params("namexxx     ").is_none());
    }

    #[test]
    fn test_parse_position_params() {
        use super::parse_position_params;
        assert_eq!(parse_position_params("startpos  ").unwrap().fen,
                   "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 1");
        assert_eq!(parse_position_params("startpos ").unwrap().moves.len(), 0);
        assert_eq!(parse_position_params("startpos   moves  ").unwrap().moves.len(),
                   0);
        assert_eq!(parse_position_params("startpos   moves   e2e4   d2d4 ").unwrap().moves.len(),
                   2);
        assert_eq!(parse_position_params("fen xxx moves e2e4").unwrap().moves.len(),
                   1);
        assert_eq!(parse_position_params("fen xxx moves e2e4").unwrap().fen,
                   "xxx".to_string());
        assert_eq!(parse_position_params("fen   xxx    moves e2e4").unwrap().fen,
                   "xxx".to_string());
        assert_eq!(parse_position_params("fen   xxx    moves").unwrap().fen,
                   "xxx".to_string());
        assert_eq!(parse_position_params("fen   xxx   ").unwrap().fen,
                   "xxx".to_string());
    }
}
<|endoftext|>


<commit_before>use core::str::StrPrelude;

use io;

const BACKSPACE: u8 = 0x08;
const TAB: u8 = 0x09;
const NEWLINE: u8 = 0x0A;
const CR: u8 = 0x0D;
const WHITESPACE: u8 = 0x20;

const VGA_ADDRESS: int = 0xB8000;
const VGA_HEIGHT: u16 = 25;
const VGA_WIDTH: u16 = 80;

enum Color {
    Black = 0,
    Blue = 1,
    Green = 2,
    Cyan = 3,
    Red = 4,
    Pink = 5,
    Brown = 6,
    LightGray = 7,
    DarkGray = 8,
    LightBlue = 9,
    LightGreen = 10,
    LightCyan = 11,
    LightRed = 12,
    LightPink = 13,
    Yellow = 14,
    White = 15,
}

static mut VGA: VGA = VGA { height: VGA_HEIGHT, width: VGA_WIDTH, x: 0, y: 0 };

struct VGA {
    height: u16,
    width: u16,
    x: u16,
    y: u16,
}

impl VGA {

    fn new() -> VGA {
        VGA {height: VGA_HEIGHT, width: VGA_WIDTH, x: 0, y: 0}
    }

    fn back(&mut self) {

        if self.x > 0 {
            self.x -= 1;

            let offset = self.offset();
            self.put(offset, Black as u16, White as u16, WHITESPACE);
        }
    }

    fn forward(&mut self) {
        self.x += 1;

        if self.x >= self.width {
            self.newline();
        }
    }

    fn cr(&mut self) {
        self.x = 0;
    }

    fn newline(&mut self) {
        self.x = 0;
        self.y += 1;
    }

    fn offset(&self) -> u16 {
        self.y * self.width + self.x
    }

    fn reset(&mut self) {
        self.x = 0;
        self.y = 0;
        self.mov();
    }

    // handle a tab by increasing the cursor's X, but only to a point
    // where it is divisible by 8
    fn tab(&mut self) {
        self.x = (self.x + 8) & !(8 - 1);
    }

    fn mov(&mut self) {
        let offset = self.offset() as u8;

        io::port::write(0x3D4, 14);  // tell the VGA board we are setting the high cursor byte
        io::port::write(0x3D5, offset >> 8); // send the high cursor byte
        io::port::write(0x3D4, 15);  // tell the VGA board we are setting the low cursor byte
        io::port::write(0x3D5, offset);
    }

    fn clear (&mut self) {
        let mut x: u16 = 0;

        loop {
            if x > 80 * 200 {
                break;
            }

            self.put(x, Black as u16, White as u16, WHITESPACE);
            x += 1;
        }
    }

    fn put(&mut self, offset: u16, background: u16, foreground: u16, character: u8) {
        let pixel: u16 = (background << 12) | (foreground << 8) | character as u16;

        unsafe {
                *((VGA_ADDRESS + offset as int * 2) as *mut u16) = pixel;
            }
    }

    fn puti(&mut self, integer: uint) {  // FIXME
        let mut integer = integer;

        loop {
            self.putc((integer % 10) as u8 + '0' as u8);
            integer = integer / 10;

            if integer == 0 {
                break;
            }
        }
    }

    fn putc(&mut self, character: u8) {

        if character == BACKSPACE {
            self.back();
        }

        else if character == TAB {
            self.tab();
        }

        else if character == NEWLINE {
            self.newline();
        }

        else if character == CR {
            self.cr();
        }

        else if character >= WHITESPACE {
            let offset = self.offset();
            self.put(offset, Black as u16, White as u16, character);
            self.forward();
        }

        self.mov();
   }

    fn puts(&mut self, string: &str) {
        for character in string.bytes() {
            self.putc(character);
        }

    }
}

pub fn clear() {
    unsafe {
        VGA.clear();
    }
}

pub fn puti(integer: uint) {
    unsafe {
        VGA.puti(integer);
    }
}

pub fn putc(character: u8) {
    unsafe {
        VGA.putc(character);
    }
}

pub fn puts(string: &str) {
    unsafe {
        VGA.puts(string);
    }
}
<commit_msg>Print hexadecimal numbers<commit_after>use core::prelude::*;

use io;

const BACKSPACE: u8 = 0x08;
const TAB: u8 = 0x09;
const NEWLINE: u8 = 0x0A;
const CR: u8 = 0x0D;
const WHITESPACE: u8 = 0x20;

const VGA_ADDRESS: int = 0xB8000;
const VGA_HEIGHT: u16 = 25;
const VGA_WIDTH: u16 = 80;

enum Color {
    Black = 0,
    Blue = 1,
    Green = 2,
    Cyan = 3,
    Red = 4,
    Pink = 5,
    Brown = 6,
    LightGray = 7,
    DarkGray = 8,
    LightBlue = 9,
    LightGreen = 10,
    LightCyan = 11,
    LightRed = 12,
    LightPink = 13,
    Yellow = 14,
    White = 15,
}

static mut VGA: VGA = VGA { height: VGA_HEIGHT, width: VGA_WIDTH, x: 0, y: 0 };

struct VGA {
    height: u16,
    width: u16,
    x: u16,
    y: u16,
}

impl VGA {

    fn new() -> VGA {
        VGA {height: VGA_HEIGHT, width: VGA_WIDTH, x: 0, y: 0}
    }

    fn back(&mut self) {

        if self.x > 0 {
            self.x -= 1;

            let offset = self.offset();
            self.put(offset, Black as u16, White as u16, WHITESPACE);
        }
    }

    fn forward(&mut self) {
        self.x += 1;

        if self.x >= self.width {
            self.newline();
        }
    }

    fn cr(&mut self) {
        self.x = 0;
    }

    fn newline(&mut self) {
        self.x = 0;
        self.y += 1;
    }

    fn offset(&self) -> u16 {
        self.y * self.width + self.x
    }

    fn reset(&mut self) {
        self.x = 0;
        self.y = 0;
        self.mov();
    }

    // handle a tab by increasing the cursor's X, but only to a point
    // where it is divisible by 8
    fn tab(&mut self) {
        self.x = (self.x + 8) & !(8 - 1);
    }

    fn mov(&mut self) {
        let offset = self.offset() as u8;

        io::port::write(0x3D4, 14);  // tell the VGA board we are setting the high cursor byte
        io::port::write(0x3D5, offset >> 8); // send the high cursor byte
        io::port::write(0x3D4, 15);  // tell the VGA board we are setting the low cursor byte
        io::port::write(0x3D5, offset);
    }

    fn clear (&mut self) {
        let mut x: u16 = 0;

        loop {
            if x > 80 * 200 {
                break;
            }

            self.put(x, Black as u16, White as u16, WHITESPACE);
            x += 1;
        }
    }

    fn put(&mut self, offset: u16, background: u16, foreground: u16, character: u8) {
        let pixel: u16 = (background << 12) | (foreground << 8) | character as u16;

        unsafe {
                *((VGA_ADDRESS + offset as int * 2) as *mut u16) = pixel;
            }
    }

    fn puti(&mut self, integer: uint) {  // FIXME
        let mut integer = integer;

        loop {
            self.putc((integer % 10) as u8 + '0' as u8);
            integer = integer / 10;

            if integer == 0 {
                break;
            }
        }
    }

    fn putc(&mut self, character: u8) {

        if character == BACKSPACE {
            self.back();
        }

        else if character == TAB {
            self.tab();
        }

        else if character == NEWLINE {
            self.newline();
        }

        else if character == CR {
            self.cr();
        }

        else if character >= WHITESPACE {
            let offset = self.offset();
            self.put(offset, Black as u16, White as u16, character);
            self.forward();
        }

        self.mov();
   }

    fn puth(&mut self, integer: uint) {
        self.puts("0x");

        let mut nibbles = 1u;

        while (integer >> nibbles * 4) > 0 {
                nibbles += 1
        }

        for i in range(0, nibbles) {
            let nibble = ((integer >> (nibbles - i - 1) * 4) & 0xF) as u8;
            let character = if nibble < 10 { '0' as u8 + nibble } else { 'a' as u8 + nibble - 10 };
            self.putc(character);
        }
    }

    fn puts(&mut self, string: &str) {
        for character in string.bytes() {
            self.putc(character);
        }

    }
}

pub fn clear() {
    unsafe {
        VGA.clear();
    }
}

pub fn puth(integer: uint) {
    unsafe {
        VGA.puth(integer);
    }
}

pub fn puti(integer: uint) {
    unsafe {
        VGA.puti(integer);
    }
}

pub fn putc(character: u8) {
    unsafe {
        VGA.putc(character);
    }
}

pub fn puts(string: &str) {
    unsafe {
        VGA.puts(string);
    }
}
<|endoftext|>


<commit_before>use super::errors::{AttesterSlashingInvalid as Invalid, AttesterSlashingValidationError as Error};
use super::verify_indexed_attestation::verify_indexed_attestation;
use types::*;

/// Indicates if an `AttesterSlashing` is valid to be included in a block in the current epoch of the given
/// state.
///
/// Returns `Ok(())` if the `AttesterSlashing` is valid, otherwise indicates the reason for invalidity.
///
/// Spec v0.5.1
pub fn verify_attester_slashing<T: EthSpec>(
    state: &BeaconState<T>,
    attester_slashing: &AttesterSlashing,
    should_verify_indexed_attestations: bool,
    spec: &ChainSpec,
) -> Result<(), Error> {
    let indexed_attestation_1 = &attester_slashing.indexed_attestation_1;
    let indexed_attestation_2 = &attester_slashing.indexed_attestation_2;

    verify!(
        indexed_attestation_1.data != indexed_attestation_2.data,
        Invalid::AttestationDataIdentical
    );
    verify!(
        indexed_attestation_1.is_double_vote(indexed_attestation_2, spec)
            | indexed_attestation_1.is_surround_vote(indexed_attestation_2, spec),
        Invalid::NotSlashable
    );

    if should_verify_indexed_attestations {
        verify_indexed_attestation(state, &indexed_attestation_1, spec)
            .map_err(|e| Error::Invalid(Invalid::IndexedAttestation1Invalid(e.into())))?;
        verify_indexed_attestation(state, &indexed_attestation_2, spec)
            .map_err(|e| Error::Invalid(Invalid::IndexedAttestation2Invalid(e.into())))?;
    }

    Ok(())
}

/// For a given attester slashing, return the indices able to be slashed.
///
/// Returns Ok(indices) if `indices.len() > 0`.
///
/// Spec v0.5.1
pub fn gather_attester_slashing_indices<T: EthSpec>(
    state: &BeaconState<T>,
    attester_slashing: &AttesterSlashing,
    spec: &ChainSpec,
) -> Result<Vec<u64>, Error> {
    gather_attester_slashing_indices_modular(
        state,
        attester_slashing,
        |_, validator| validator.slashed,
        spec,
    )
}

/// Same as `gather_attester_slashing_indices` but allows the caller to specify the criteria
/// for determining whether a given validator should be considered slashed.
pub fn gather_attester_slashing_indices_modular<F, T: EthSpec>(
    state: &BeaconState<T>,
    attester_slashing: &AttesterSlashing,
    is_slashed: F,
    spec: &ChainSpec,
) -> Result<Vec<u64>, Error>
where
    F: Fn(u64, &Validator) -> bool,
{
    let indexed_attestation_1 = &attester_slashing.indexed_attestation_1;
    let indexed_attestation_2 = &attester_slashing.indexed_attestation_2;

    let mut indexed_indices = Vec::with_capacity(spec.max_indices_per_indexed_vote);
    for i in &indexed_attestation_1.validator_indices {
        let validator = state
            .validator_registry
            .get(*i as usize)
            .ok_or_else(|| Error::Invalid(Invalid::UnknownValidator(*i)))?;

        if indexed_attestation_2.validator_indices.contains(&i) & !is_slashed(*i, validator) {
            // TODO: verify that we should reject any indexed attestation which includes a
            // withdrawn validator. PH has asked the question on gitter, awaiting response.
            verify!(
                validator.withdrawable_epoch > state.slot.epoch(spec.slots_per_epoch),
                Invalid::ValidatorAlreadyWithdrawn(*i)
            );

            indexed_indices.push(*i);
        }
    }

    verify!(!indexed_indices.is_empty(), Invalid::NoSlashableIndices);

    indexed_indices.shrink_to_fit();

    Ok(indexed_indices)
}
<commit_msg>spec v0.6.1: verify attester slashing<commit_after>use super::errors::{AttesterSlashingInvalid as Invalid, AttesterSlashingValidationError as Error};
use super::verify_indexed_attestation::verify_indexed_attestation;
use types::*;

/// Indicates if an `AttesterSlashing` is valid to be included in a block in the current epoch of the given
/// state.
///
/// Returns `Ok(())` if the `AttesterSlashing` is valid, otherwise indicates the reason for invalidity.
///
/// Spec v0.6.1
pub fn verify_attester_slashing<T: EthSpec>(
    state: &BeaconState<T>,
    attester_slashing: &AttesterSlashing,
    should_verify_indexed_attestations: bool,
    spec: &ChainSpec,
) -> Result<(), Error> {
    let attestation_1 = &attester_slashing.attestation_1;
    let attestation_2 = &attester_slashing.attestation_2;

    // Spec: is_slashable_attestation_data
    verify!(
        attestation_1.is_double_vote(attestation_2, spec)
            || attestation_1.is_surround_vote(attestation_2, spec),
        Invalid::NotSlashable
    );

    if should_verify_indexed_attestations {
        verify_indexed_attestation(state, &attestation_1, spec)
            .map_err(|e| Error::Invalid(Invalid::IndexedAttestation1Invalid(e.into())))?;
        verify_indexed_attestation(state, &attestation_2, spec)
            .map_err(|e| Error::Invalid(Invalid::IndexedAttestation2Invalid(e.into())))?;
    }

    Ok(())
}

/// For a given attester slashing, return the indices able to be slashed.
///
/// Returns Ok(indices) if `indices.len() > 0`.
///
/// Spec v0.6.1
pub fn get_slashable_indices<T: EthSpec>(
    state: &BeaconState<T>,
    attester_slashing: &AttesterSlashing,
    spec: &ChainSpec,
) -> Result<Vec<u64>, Error> {
    gather_attester_slashing_indices_modular(
        state,
        attester_slashing,
        |_, validator| validator.is_slashable_at(state.current_epoch()),
        spec,
    )
}

/// Same as `gather_attester_slashing_indices` but allows the caller to specify the criteria
/// for determining whether a given validator should be considered slashable.
pub fn get_slashable_indices_modular<F, T: EthSpec>(
    state: &BeaconState<T>,
    attester_slashing: &AttesterSlashing,
    is_slashable: F,
    spec: &ChainSpec,
) -> Result<Vec<u64>, Error>
where
    F: Fn(u64, &Validator) -> bool,
{
    let attestation_1 = &attester_slashing.attestation_1;
    let attestation_2 = &attester_slashing.attestation_2;

    let mut attesting_indices_1 = HashSet::new();
    attesting_indices_1.extend(attestation_1.custody_bit_0_indices.clone());
    attesting_indices_1.extend(attestation_1.custody_bit_1_indices.clone());

    let mut attesting_indices_2 = HashSet::new();
    attesting_indices_2.extend(attestation_2.custody_bit_0_indices.clone());
    attesting_indices_2.extend(attestation_2.custody_bit_1_indices.clone());

    let mut slashable_indices = vec![];

    for index in &attesting_indices_1 & &attesting_indices_2 {
        let validator = state
            .validator_registry
            .get(index as usize)
            .ok_or_else(|| Error::Invalid(Invalid::UnknownValidator(index)))?;

        if is_slashable(index, validator) {
            slashable_indices.push(index);
        }
    }

    verify!(!slashable_indices.is_empty(), Invalid::NoSlashableIndices);

    Ok(slashable_indices)
}
<|endoftext|>


<commit_before>/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */

use aws_config::meta::region::RegionProviderChain;
use aws_sdk_globalaccelerator::{Client, Error, Region, PKG_VERSION};
use structopt::StructOpt;
use tokio_stream::StreamExt;

#[derive(Debug, StructOpt)]
struct Opt {
    /// Whether to display additional information.
    #[structopt(short, long)]
    verbose: bool,
}

// List accelerator names and ARNs.
// snippet-start:[globalaccelerator.rust.globalaccelerator-helloworld]
async fn show_accelerators(client: &Client) -> Result<(), Error> {
    let mut paginator = client.list_accelerators().into_paginator().send();

    while let Some(page) = paginator.try_next().await? {
        for accelerator in page.accelerators().unwrap_or_default().iter() {
            let accelerator_arn = accelerator.name().unwrap_or_default();
            let accelerator_name = accelerator.accelerator_arn().unwrap_or_default();

            println!("Accelerator Name : {}", accelerator_name);
            println!("Accelerator ARN : {}", accelerator_arn);
            println!();
        }
    }

    Ok(())
}
// snippet-end:[globalaccelerator.rust.globalaccelerator-helloworld]

/// Lists your AWS Global Accelerator accelerator names and ARNs in the Region.
/// # Arguments
/// * `[-v]` - Whether to display additional information.
#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt::init();

    let Opt { verbose } = Opt::from_args();

    // Global Accelerator is a global service with its API in us-west-2
    let region_provider = RegionProviderChain::first_try(Region::new("us-west-2"));
    println!();

    if verbose {
        println!("globalaccelerator client version: {}", PKG_VERSION);
        println!(
            "Region:                   {}",
            region_provider.region().await.unwrap().as_ref()
        );
        println!();
    }

    let shared_config = aws_config::from_env().region(region_provider).load().await;
    let client = Client::new(&shared_config);

    show_accelerators(&client).await
}
<commit_msg>Add an intro output<commit_after>/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */

use aws_config::meta::region::RegionProviderChain;
use aws_sdk_globalaccelerator::{Client, Error, Region, PKG_VERSION};
use structopt::StructOpt;
use tokio_stream::StreamExt;

#[derive(Debug, StructOpt)]
struct Opt {
    /// Whether to display additional information.
    #[structopt(short, long)]
    verbose: bool,
}

// List accelerator names and ARNs.
// snippet-start:[globalaccelerator.rust.globalaccelerator-helloworld]
async fn show_accelerators(client: &Client) -> Result<(), Error> {
    println!("Welcome to the AWS Rust SDK Global Accelerator example!");
    println!();
    let mut paginator = client.list_accelerators().into_paginator().send();

    while let Some(page) = paginator.try_next().await? {
        for accelerator in page.accelerators().unwrap_or_default().iter() {
            let accelerator_arn = accelerator.name().unwrap_or_default();
            let accelerator_name = accelerator.accelerator_arn().unwrap_or_default();

            println!("Accelerator Name : {}", accelerator_name);
            println!("Accelerator ARN : {}", accelerator_arn);
            println!();
        }
    }

    Ok(())
}
// snippet-end:[globalaccelerator.rust.globalaccelerator-helloworld]

/// Lists your AWS Global Accelerator accelerator names and ARNs in the Region.
/// # Arguments
/// * `[-v]` - Whether to display additional information.
#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt::init();

    let Opt { verbose } = Opt::from_args();

    // Global Accelerator is a global service with its API in us-west-2
    let region_provider = RegionProviderChain::first_try(Region::new("us-west-2"));
    println!();

    if verbose {
        println!("globalaccelerator client version: {}", PKG_VERSION);
        println!(
            "Region:                   {}",
            region_provider.region().await.unwrap().as_ref()
        );
        println!();
    }

    let shared_config = aws_config::from_env().region(region_provider).load().await;
    let client = Client::new(&shared_config);

    show_accelerators(&client).await
}
<|endoftext|>


<commit_before>use crate::clang::{Clang, Node};
use crate::syntax::attrs::OtherAttrs;
use crate::syntax::cfg::CfgExpr;
use crate::syntax::namespace::Namespace;
use crate::syntax::report::Errors;
use crate::syntax::{Api, Discriminant, Doc, Enum, EnumRepr, ForeignName, Pair, Variant};
use flate2::write::GzDecoder;
use memmap::Mmap;
use proc_macro2::{Delimiter, Group, Ident, TokenStream};
use quote::{format_ident, quote, quote_spanned};
use std::env;
use std::fmt::{self, Display};
use std::fs::File;
use std::io::Write;
use std::path::PathBuf;
use std::str::FromStr;
use syn::{parse_quote, Path};

const CXX_CLANG_AST: &str = "CXX_CLANG_AST";

pub fn load(cx: &mut Errors, apis: &mut [Api]) {
    let ref mut variants_from_header = Vec::new();
    for api in apis {
        if let Api::Enum(enm) = api {
            if enm.variants_from_header {
                if enm.variants.is_empty() {
                    variants_from_header.push(enm);
                } else {
                    let span = span_for_enum_error(enm);
                    cx.error(
                        span,
                        "enum with #![variants_from_header] must be written with no explicit variants",
                    );
                }
            }
        }
    }

    let span = match variants_from_header.get(0) {
        None => return,
        Some(enm) => enm.variants_from_header_attr.clone().unwrap(),
    };

    let ast_dump_path = match env::var_os(CXX_CLANG_AST) {
        Some(ast_dump_path) => PathBuf::from(ast_dump_path),
        None => {
            let msg = format!(
                "environment variable ${} has not been provided",
                CXX_CLANG_AST,
            );
            return cx.error(span, msg);
        }
    };

    let memmap = File::open(&ast_dump_path).and_then(|file| unsafe { Mmap::map(&file) });
    let mut gunzipped;
    let ast_dump_bytes = match match memmap {
        Ok(ref memmap) => {
            let is_gzipped = memmap.get(..2) == Some(b"\x1f\x8b");
            if is_gzipped {
                gunzipped = Vec::new();
                let decode_result = GzDecoder::new(&mut gunzipped).write_all(&memmap);
                decode_result.map(|_| gunzipped.as_slice())
            } else {
                Ok(&memmap as &[u8])
            }
        }
        Err(error) => Err(error),
    } {
        Ok(bytes) => bytes,
        Err(error) => {
            let msg = format!("failed to read {}: {}", ast_dump_path.display(), error);
            return cx.error(span, msg);
        }
    };

    let ref root: Node = match serde_json::from_slice(ast_dump_bytes) {
        Ok(root) => root,
        Err(error) => {
            let msg = format!("failed to read {}: {}", ast_dump_path.display(), error);
            return cx.error(span, msg);
        }
    };

    let ref mut namespace = Vec::new();
    traverse(cx, root, namespace, variants_from_header, None);

    for enm in variants_from_header {
        if enm.variants.is_empty() {
            let span = &enm.variants_from_header_attr;
            let name = CxxName(&enm.name);
            let msg = format!("failed to find any C++ definition of enum {}", name);
            cx.error(span, msg);
        }
    }
}

fn traverse<'a>(
    cx: &mut Errors,
    node: &'a Node,
    namespace: &mut Vec<&'a str>,
    variants_from_header: &mut [&mut Enum],
    mut idx: Option<usize>,
) {
    match &node.kind {
        Clang::NamespaceDecl(decl) => {
            let name = match &decl.name {
                Some(name) => name,
                // Can ignore enums inside an anonymous namespace.
                None => return,
            };
            namespace.push(name);
            idx = None;
        }
        Clang::EnumDecl(decl) => {
            let name = match &decl.name {
                Some(name) => name,
                None => return,
            };
            idx = None;
            for (i, enm) in variants_from_header.iter_mut().enumerate() {
                if enm.name.cxx == **name && enm.name.namespace.iter().eq(&*namespace) {
                    if !enm.variants.is_empty() {
                        let span = &enm.variants_from_header_attr;
                        let qual_name = CxxName(&enm.name);
                        let msg = format!("found multiple C++ definitions of enum {}", qual_name);
                        cx.error(span, msg);
                        return;
                    }
                    let fixed_underlying_type = match &decl.fixed_underlying_type {
                        Some(fixed_underlying_type) => fixed_underlying_type,
                        None => {
                            let span = &enm.variants_from_header_attr;
                            let name = &enm.name.cxx;
                            let qual_name = CxxName(&enm.name);
                            let msg = format!(
                                "implicit implementation-defined repr for enum {} is not supported yet; consider changing its C++ definition to `enum {}: int {{...}}",
                                qual_name, name,
                            );
                            cx.error(span, msg);
                            return;
                        }
                    };
                    let repr = translate_qual_type(
                        cx,
                        &enm,
                        fixed_underlying_type
                            .desugared_qual_type
                            .as_ref()
                            .unwrap_or(&fixed_underlying_type.qual_type),
                    );
                    enm.repr = EnumRepr::Foreign { rust_type: repr };
                    idx = Some(i);
                    break;
                }
            }
            if idx.is_none() {
                return;
            }
        }
        Clang::EnumConstantDecl(decl) => {
            if let Some(idx) = idx {
                let enm = &mut *variants_from_header[idx];
                let span = enm
                    .variants_from_header_attr
                    .as_ref()
                    .unwrap()
                    .path
                    .get_ident()
                    .unwrap()
                    .span();
                let cxx_name = match ForeignName::parse(&decl.name, span) {
                    Ok(foreign_name) => foreign_name,
                    Err(_) => {
                        let span = &enm.variants_from_header_attr;
                        let msg = format!("unsupported C++ variant name: {}", decl.name);
                        return cx.error(span, msg);
                    }
                };
                let rust_name: Ident = match syn::parse_str(&decl.name) {
                    Ok(ident) => ident,
                    Err(_) => format_ident!("__Variant{}", enm.variants.len()),
                };
                let discriminant = match discriminant_value(&node.inner) {
                    ParsedDiscriminant::Constant(discriminant) => discriminant,
                    ParsedDiscriminant::Successor => match enm.variants.last() {
                        None => Discriminant::zero(),
                        Some(last) => match last.discriminant.checked_succ() {
                            Some(discriminant) => discriminant,
                            None => {
                                let span = &enm.variants_from_header_attr;
                                let msg = format!(
                                    "overflow processing discriminant value for variant: {}",
                                    decl.name,
                                );
                                return cx.error(span, msg);
                            }
                        },
                    },
                    ParsedDiscriminant::Fail => {
                        let span = &enm.variants_from_header_attr;
                        let msg = format!(
                            "failed to obtain discriminant value for variant: {}",
                            decl.name,
                        );
                        cx.error(span, msg);
                        Discriminant::zero()
                    }
                };
                enm.variants.push(Variant {
                    cfg: CfgExpr::Unconditional,
                    doc: Doc::new(),
                    attrs: OtherAttrs::none(),
                    name: Pair {
                        namespace: Namespace::ROOT,
                        cxx: cxx_name,
                        rust: rust_name,
                    },
                    discriminant,
                    expr: None,
                });
            }
        }
        _ => {}
    }
    for inner in &node.inner {
        traverse(cx, inner, namespace, variants_from_header, idx);
    }
    if let Clang::NamespaceDecl(_) = &node.kind {
        let _ = namespace.pop().unwrap();
    }
}

fn translate_qual_type(cx: &mut Errors, enm: &Enum, qual_type: &str) -> Path {
    let rust_std_name = match qual_type {
        "char" => "c_char",
        "int" => "c_int",
        "long" => "c_long",
        "long long" => "c_longlong",
        "signed char" => "c_schar",
        "short" => "c_short",
        "unsigned char" => "c_uchar",
        "unsigned int" => "c_uint",
        "unsigned long" => "c_ulong",
        "unsigned long long" => "c_ulonglong",
        "unsigned short" => "c_ushort",
        unsupported => {
            let span = &enm.variants_from_header_attr;
            let qual_name = CxxName(&enm.name);
            let msg = format!(
                "unsupported underlying type for {}: {}",
                qual_name, unsupported,
            );
            cx.error(span, msg);
            "c_int"
        }
    };
    let span = enm
        .variants_from_header_attr
        .as_ref()
        .unwrap()
        .path
        .get_ident()
        .unwrap()
        .span();
    let ident = Ident::new(rust_std_name, span);
    let path = quote_spanned!(span=> ::std::os::raw::#ident);
    parse_quote!(#path)
}

enum ParsedDiscriminant {
    Constant(Discriminant),
    Successor,
    Fail,
}

fn discriminant_value(mut clang: &[Node]) -> ParsedDiscriminant {
    if clang.is_empty() {
        // No discriminant expression provided; use successor of previous
        // descriminant.
        return ParsedDiscriminant::Successor;
    }

    loop {
        if clang.len() != 1 {
            return ParsedDiscriminant::Fail;
        }

        let node = &clang[0];
        match &node.kind {
            Clang::ImplicitCastExpr => clang = &node.inner,
            Clang::ConstantExpr(expr) => match Discriminant::from_str(&expr.value) {
                Ok(discriminant) => return ParsedDiscriminant::Constant(discriminant),
                Err(_) => return ParsedDiscriminant::Fail,
            },
            _ => return ParsedDiscriminant::Fail,
        }
    }
}

fn span_for_enum_error(enm: &Enum) -> TokenStream {
    let enum_token = enm.enum_token;
    let mut brace_token = Group::new(Delimiter::Brace, TokenStream::new());
    brace_token.set_span(enm.brace_token.span);
    quote!(#enum_token #brace_token)
}

struct CxxName<'a>(&'a Pair);

impl<'a> Display for CxxName<'a> {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        for namespace in &self.0.namespace {
            write!(formatter, "{}::", namespace)?;
        }
        write!(formatter, "{}", self.0.cxx)
    }
}
<commit_msg>clippy fixes<commit_after>use crate::clang::{Clang, Node};
use crate::syntax::attrs::OtherAttrs;
use crate::syntax::cfg::CfgExpr;
use crate::syntax::namespace::Namespace;
use crate::syntax::report::Errors;
use crate::syntax::{Api, Discriminant, Doc, Enum, EnumRepr, ForeignName, Pair, Variant};
use flate2::write::GzDecoder;
use memmap::Mmap;
use proc_macro2::{Delimiter, Group, Ident, TokenStream};
use quote::{format_ident, quote, quote_spanned};
use std::env;
use std::fmt::{self, Display};
use std::fs::File;
use std::io::Write;
use std::path::PathBuf;
use std::str::FromStr;
use syn::{parse_quote, Path};

const CXX_CLANG_AST: &str = "CXX_CLANG_AST";

pub fn load(cx: &mut Errors, apis: &mut [Api]) {
    let ref mut variants_from_header = Vec::new();
    for api in apis {
        if let Api::Enum(enm) = api {
            if enm.variants_from_header {
                if enm.variants.is_empty() {
                    variants_from_header.push(enm);
                } else {
                    let span = span_for_enum_error(enm);
                    cx.error(
                        span,
                        "enum with #![variants_from_header] must be written with no explicit variants",
                    );
                }
            }
        }
    }

    let span = match variants_from_header.get(0) {
        None => return,
        Some(enm) => enm.variants_from_header_attr.clone().unwrap(),
    };

    let ast_dump_path = match env::var_os(CXX_CLANG_AST) {
        Some(ast_dump_path) => PathBuf::from(ast_dump_path),
        None => {
            let msg = format!(
                "environment variable ${} has not been provided",
                CXX_CLANG_AST,
            );
            return cx.error(span, msg);
        }
    };

    let memmap = File::open(&ast_dump_path).and_then(|file| unsafe { Mmap::map(&file) });
    let mut gunzipped;
    let ast_dump_bytes = match match memmap {
        Ok(ref memmap) => {
            let is_gzipped = memmap.get(..2) == Some(b"\x1f\x8b");
            if is_gzipped {
                gunzipped = Vec::new();
                let decode_result = GzDecoder::new(&mut gunzipped).write_all(memmap);
                decode_result.map(|_| gunzipped.as_slice())
            } else {
                Ok(memmap as &[u8])
            }
        }
        Err(error) => Err(error),
    } {
        Ok(bytes) => bytes,
        Err(error) => {
            let msg = format!("failed to read {}: {}", ast_dump_path.display(), error);
            return cx.error(span, msg);
        }
    };

    let ref root: Node = match serde_json::from_slice(ast_dump_bytes) {
        Ok(root) => root,
        Err(error) => {
            let msg = format!("failed to read {}: {}", ast_dump_path.display(), error);
            return cx.error(span, msg);
        }
    };

    let ref mut namespace = Vec::new();
    traverse(cx, root, namespace, variants_from_header, None);

    for enm in variants_from_header {
        if enm.variants.is_empty() {
            let span = &enm.variants_from_header_attr;
            let name = CxxName(&enm.name);
            let msg = format!("failed to find any C++ definition of enum {}", name);
            cx.error(span, msg);
        }
    }
}

fn traverse<'a>(
    cx: &mut Errors,
    node: &'a Node,
    namespace: &mut Vec<&'a str>,
    variants_from_header: &mut [&mut Enum],
    mut idx: Option<usize>,
) {
    match &node.kind {
        Clang::NamespaceDecl(decl) => {
            let name = match &decl.name {
                Some(name) => name,
                // Can ignore enums inside an anonymous namespace.
                None => return,
            };
            namespace.push(name);
            idx = None;
        }
        Clang::EnumDecl(decl) => {
            let name = match &decl.name {
                Some(name) => name,
                None => return,
            };
            idx = None;
            for (i, enm) in variants_from_header.iter_mut().enumerate() {
                if enm.name.cxx == **name && enm.name.namespace.iter().eq(&*namespace) {
                    if !enm.variants.is_empty() {
                        let span = &enm.variants_from_header_attr;
                        let qual_name = CxxName(&enm.name);
                        let msg = format!("found multiple C++ definitions of enum {}", qual_name);
                        cx.error(span, msg);
                        return;
                    }
                    let fixed_underlying_type = match &decl.fixed_underlying_type {
                        Some(fixed_underlying_type) => fixed_underlying_type,
                        None => {
                            let span = &enm.variants_from_header_attr;
                            let name = &enm.name.cxx;
                            let qual_name = CxxName(&enm.name);
                            let msg = format!(
                                "implicit implementation-defined repr for enum {} is not supported yet; consider changing its C++ definition to `enum {}: int {{...}}",
                                qual_name, name,
                            );
                            cx.error(span, msg);
                            return;
                        }
                    };
                    let repr = translate_qual_type(
                        cx,
                        enm,
                        fixed_underlying_type
                            .desugared_qual_type
                            .as_ref()
                            .unwrap_or(&fixed_underlying_type.qual_type),
                    );
                    enm.repr = EnumRepr::Foreign { rust_type: repr };
                    idx = Some(i);
                    break;
                }
            }
            if idx.is_none() {
                return;
            }
        }
        Clang::EnumConstantDecl(decl) => {
            if let Some(idx) = idx {
                let enm = &mut *variants_from_header[idx];
                let span = enm
                    .variants_from_header_attr
                    .as_ref()
                    .unwrap()
                    .path
                    .get_ident()
                    .unwrap()
                    .span();
                let cxx_name = match ForeignName::parse(&decl.name, span) {
                    Ok(foreign_name) => foreign_name,
                    Err(_) => {
                        let span = &enm.variants_from_header_attr;
                        let msg = format!("unsupported C++ variant name: {}", decl.name);
                        return cx.error(span, msg);
                    }
                };
                let rust_name: Ident = match syn::parse_str(&decl.name) {
                    Ok(ident) => ident,
                    Err(_) => format_ident!("__Variant{}", enm.variants.len()),
                };
                let discriminant = match discriminant_value(&node.inner) {
                    ParsedDiscriminant::Constant(discriminant) => discriminant,
                    ParsedDiscriminant::Successor => match enm.variants.last() {
                        None => Discriminant::zero(),
                        Some(last) => match last.discriminant.checked_succ() {
                            Some(discriminant) => discriminant,
                            None => {
                                let span = &enm.variants_from_header_attr;
                                let msg = format!(
                                    "overflow processing discriminant value for variant: {}",
                                    decl.name,
                                );
                                return cx.error(span, msg);
                            }
                        },
                    },
                    ParsedDiscriminant::Fail => {
                        let span = &enm.variants_from_header_attr;
                        let msg = format!(
                            "failed to obtain discriminant value for variant: {}",
                            decl.name,
                        );
                        cx.error(span, msg);
                        Discriminant::zero()
                    }
                };
                enm.variants.push(Variant {
                    cfg: CfgExpr::Unconditional,
                    doc: Doc::new(),
                    attrs: OtherAttrs::none(),
                    name: Pair {
                        namespace: Namespace::ROOT,
                        cxx: cxx_name,
                        rust: rust_name,
                    },
                    discriminant,
                    expr: None,
                });
            }
        }
        _ => {}
    }
    for inner in &node.inner {
        traverse(cx, inner, namespace, variants_from_header, idx);
    }
    if let Clang::NamespaceDecl(_) = &node.kind {
        let _ = namespace.pop().unwrap();
    }
}

fn translate_qual_type(cx: &mut Errors, enm: &Enum, qual_type: &str) -> Path {
    let rust_std_name = match qual_type {
        "char" => "c_char",
        "int" => "c_int",
        "long" => "c_long",
        "long long" => "c_longlong",
        "signed char" => "c_schar",
        "short" => "c_short",
        "unsigned char" => "c_uchar",
        "unsigned int" => "c_uint",
        "unsigned long" => "c_ulong",
        "unsigned long long" => "c_ulonglong",
        "unsigned short" => "c_ushort",
        unsupported => {
            let span = &enm.variants_from_header_attr;
            let qual_name = CxxName(&enm.name);
            let msg = format!(
                "unsupported underlying type for {}: {}",
                qual_name, unsupported,
            );
            cx.error(span, msg);
            "c_int"
        }
    };
    let span = enm
        .variants_from_header_attr
        .as_ref()
        .unwrap()
        .path
        .get_ident()
        .unwrap()
        .span();
    let ident = Ident::new(rust_std_name, span);
    let path = quote_spanned!(span=> ::std::os::raw::#ident);
    parse_quote!(#path)
}

enum ParsedDiscriminant {
    Constant(Discriminant),
    Successor,
    Fail,
}

fn discriminant_value(mut clang: &[Node]) -> ParsedDiscriminant {
    if clang.is_empty() {
        // No discriminant expression provided; use successor of previous
        // descriminant.
        return ParsedDiscriminant::Successor;
    }

    loop {
        if clang.len() != 1 {
            return ParsedDiscriminant::Fail;
        }

        let node = &clang[0];
        match &node.kind {
            Clang::ImplicitCastExpr => clang = &node.inner,
            Clang::ConstantExpr(expr) => match Discriminant::from_str(&expr.value) {
                Ok(discriminant) => return ParsedDiscriminant::Constant(discriminant),
                Err(_) => return ParsedDiscriminant::Fail,
            },
            _ => return ParsedDiscriminant::Fail,
        }
    }
}

fn span_for_enum_error(enm: &Enum) -> TokenStream {
    let enum_token = enm.enum_token;
    let mut brace_token = Group::new(Delimiter::Brace, TokenStream::new());
    brace_token.set_span(enm.brace_token.span);
    quote!(#enum_token #brace_token)
}

struct CxxName<'a>(&'a Pair);

impl<'a> Display for CxxName<'a> {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        for namespace in &self.0.namespace {
            write!(formatter, "{}::", namespace)?;
        }
        write!(formatter, "{}", self.0.cxx)
    }
}
<|endoftext|>