<commit_before>//! Serde `Deserializer` module

use crate::{
    de::escape::EscapedDeserializer,
    de::{DeEvent, Deserializer, XmlRead, INNER_VALUE, UNFLATTEN_PREFIX},
    errors::serialize::DeError,
    events::attributes::IterState,
    events::BytesStart,
};
use serde::de::{self, DeserializeSeed, IntoDeserializer};
use std::borrow::Cow;
use std::ops::Range;

/// Defines a source that should be used to deserialize a value in the next call
/// to [`next_value_seed()`](de::MapAccess::next_value_seed)
#[derive(Debug, PartialEq)]
enum ValueSource {
    /// Source are not specified, because [`next_key_seed()`] not yet called.
    /// This is an initial state and state after deserializing value
    /// (after call of [`next_value_seed()`]).
    ///
    /// Attempt to call [`next_value_seed()`] while accessor in this state would
    /// return a [`DeError::KeyNotRead`] error.
    ///
    /// [`next_key_seed()`]: de::MapAccess::next_key_seed
    /// [`next_value_seed()`]: de::MapAccess::next_value_seed
    Unknown,
    /// `next_key_seed` checked the attributes list and find it is not exhausted yet.
    /// Next call to the `next_value_seed` will deserialize type from the attribute value
    Attribute(Range<usize>),
    /// The same as `Text`
    Nested,
    /// Value should be deserialized from the text content of the XML node, which
    /// represented or by an ordinary text node, or by a CDATA node:
    ///
    /// ```xml
    /// <...>text content for field value<...>
    /// ```
    /// ```xml
    /// <any-tag>
    ///     <key><![CDATA[cdata content]]></key>
    /// <!--              ^^^^^^^^^^^^^ - this will be used to deserialize a map value -->
    /// </any-tag>
    /// ```
    Text,
}

/// A deserializer for `Attributes`
pub(crate) struct MapAccess<'de, 'a, R>
where
    R: XmlRead<'de>,
{
    /// Tag -- owner of attributes
    start: BytesStart<'de>,
    de: &'a mut Deserializer<'de, R>,
    /// State of the iterator over attributes. Contains the next position in the
    /// inner `start` slice, from which next attribute should be parsed.
    iter: IterState,
    /// Current state of the accessor that determines what next call to API
    /// methods should return.
    source: ValueSource,
    /// list of fields yet to unflatten (defined as starting with $unflatten=)
    unflatten_fields: Vec<&'static [u8]>,
}

impl<'de, 'a, R> MapAccess<'de, 'a, R>
where
    R: XmlRead<'de>,
{
    /// Create a new MapAccess
    pub fn new(
        de: &'a mut Deserializer<'de, R>,
        start: BytesStart<'de>,
        fields: &[&'static str],
    ) -> Result<Self, DeError> {
        Ok(MapAccess {
            de,
            start,
            iter: IterState::new(0, false),
            source: ValueSource::Unknown,
            unflatten_fields: fields
                .iter()
                .filter(|f| f.starts_with(UNFLATTEN_PREFIX))
                .map(|f| f.as_bytes())
                .collect(),
        })
    }
}

impl<'de, 'a, R> de::MapAccess<'de> for MapAccess<'de, 'a, R>
where
    R: XmlRead<'de>,
{
    type Error = DeError;

    fn next_key_seed<K: DeserializeSeed<'de>>(
        &mut self,
        seed: K,
    ) -> Result<Option<K::Value>, Self::Error> {
        debug_assert_eq!(self.source, ValueSource::Unknown);

        // FIXME: There error positions counted from end of tag name - need global position
        let slice = self.start.attributes_raw();
        let decoder = self.de.reader.decoder();
        let has_value_field = self.de.has_value_field;

        if let Some(a) = self.iter.next(slice).transpose()? {
            // try getting map from attributes (key= "value")
            let (key, value) = a.into();
            self.source = ValueSource::Attribute(value.unwrap_or_default());
            seed.deserialize(EscapedDeserializer::new(
                Cow::Borrowed(&slice[key]),
                decoder,
                false,
            ))
            .map(Some)
        } else {
            // try getting from events (<key>value</key>)
            match self.de.peek()? {
                DeEvent::Text(_) | DeEvent::CData(_) => {
                    self.source = ValueSource::Text;
                    // Deserialize `key` from special attribute name which means
                    // that value should be taken from the text content of the
                    // XML node
                    seed.deserialize(INNER_VALUE.into_deserializer()).map(Some)
                }
                // Used to deserialize collections of enums, like:
                // <root>
                //
                //
                //   <C/>
                // </root>
                //
                // into
                //
                // enum Enum { A, B, С }
                // struct Root {
                //     #[serde(rename = "$value")]
                //     items: Vec<Enum>,
                // }
                // TODO: This should be handled by #[serde(flatten)]
                // See https://github.com/serde-rs/serde/issues/1905
                DeEvent::Start(_) if has_value_field => {
                    self.source = ValueSource::Text;
                    seed.deserialize(INNER_VALUE.into_deserializer()).map(Some)
                }
                DeEvent::Start(e) => {
                    let key = if let Some(p) = self
                        .unflatten_fields
                        .iter()
                        .position(|f| e.name() == &f[UNFLATTEN_PREFIX.len()..])
                    {
                        // Used to deserialize elements, like:
                        // <root>
                        //   <xxx>test</xxx>
                        // </root>
                        //
                        // into
                        //
                        // struct Root {
                        //     #[serde(rename = "$unflatten=xxx")]
                        //     xxx: String,
                        // }
                        self.source = ValueSource::Text;
                        seed.deserialize(self.unflatten_fields.remove(p).into_deserializer())
                    } else {
                        let name = Cow::Borrowed(e.local_name());
                        self.source = ValueSource::Nested;
                        seed.deserialize(EscapedDeserializer::new(name, decoder, false))
                    };
                    key.map(Some)
                }
                _ => Ok(None),
            }
        }
    }

    fn next_value_seed<K: DeserializeSeed<'de>>(
        &mut self,
        seed: K,
    ) -> Result<K::Value, Self::Error> {
        match std::mem::replace(&mut self.source, ValueSource::Unknown) {
            ValueSource::Attribute(value) => {
                let slice = self.start.attributes_raw();
                let decoder = self.de.reader.decoder();

                seed.deserialize(EscapedDeserializer::new(
                    Cow::Borrowed(&slice[value]),
                    decoder,
                    true,
                ))
            }
            ValueSource::Nested | ValueSource::Text => seed.deserialize(&mut *self.de),
            ValueSource::Unknown => Err(DeError::KeyNotRead),
        }
    }
}
<commit_msg>Use correct value source for `unflatten=` fields<commit_after>//! Serde `Deserializer` module

use crate::{
    de::escape::EscapedDeserializer,
    de::{DeEvent, Deserializer, XmlRead, INNER_VALUE, UNFLATTEN_PREFIX},
    errors::serialize::DeError,
    events::attributes::IterState,
    events::BytesStart,
};
use serde::de::{self, DeserializeSeed, IntoDeserializer};
use std::borrow::Cow;
use std::ops::Range;

/// Defines a source that should be used to deserialize a value in the next call
/// to [`next_value_seed()`](de::MapAccess::next_value_seed)
#[derive(Debug, PartialEq)]
enum ValueSource {
    /// Source are not specified, because [`next_key_seed()`] not yet called.
    /// This is an initial state and state after deserializing value
    /// (after call of [`next_value_seed()`]).
    ///
    /// Attempt to call [`next_value_seed()`] while accessor in this state would
    /// return a [`DeError::KeyNotRead`] error.
    ///
    /// [`next_key_seed()`]: de::MapAccess::next_key_seed
    /// [`next_value_seed()`]: de::MapAccess::next_value_seed
    Unknown,
    /// `next_key_seed` checked the attributes list and find it is not exhausted yet.
    /// Next call to the `next_value_seed` will deserialize type from the attribute value
    Attribute(Range<usize>),
    /// The same as `Text`
    Nested,
    /// Value should be deserialized from the text content of the XML node, which
    /// represented or by an ordinary text node, or by a CDATA node:
    ///
    /// ```xml
    /// <...>text content for field value<...>
    /// ```
    /// ```xml
    /// <any-tag>
    ///     <key><![CDATA[cdata content]]></key>
    /// <!--              ^^^^^^^^^^^^^ - this will be used to deserialize a map value -->
    /// </any-tag>
    /// ```
    Text,
}

/// A deserializer for `Attributes`
pub(crate) struct MapAccess<'de, 'a, R>
where
    R: XmlRead<'de>,
{
    /// Tag -- owner of attributes
    start: BytesStart<'de>,
    de: &'a mut Deserializer<'de, R>,
    /// State of the iterator over attributes. Contains the next position in the
    /// inner `start` slice, from which next attribute should be parsed.
    iter: IterState,
    /// Current state of the accessor that determines what next call to API
    /// methods should return.
    source: ValueSource,
    /// list of fields yet to unflatten (defined as starting with $unflatten=)
    unflatten_fields: Vec<&'static [u8]>,
}

impl<'de, 'a, R> MapAccess<'de, 'a, R>
where
    R: XmlRead<'de>,
{
    /// Create a new MapAccess
    pub fn new(
        de: &'a mut Deserializer<'de, R>,
        start: BytesStart<'de>,
        fields: &[&'static str],
    ) -> Result<Self, DeError> {
        Ok(MapAccess {
            de,
            start,
            iter: IterState::new(0, false),
            source: ValueSource::Unknown,
            unflatten_fields: fields
                .iter()
                .filter(|f| f.starts_with(UNFLATTEN_PREFIX))
                .map(|f| f.as_bytes())
                .collect(),
        })
    }
}

impl<'de, 'a, R> de::MapAccess<'de> for MapAccess<'de, 'a, R>
where
    R: XmlRead<'de>,
{
    type Error = DeError;

    fn next_key_seed<K: DeserializeSeed<'de>>(
        &mut self,
        seed: K,
    ) -> Result<Option<K::Value>, Self::Error> {
        debug_assert_eq!(self.source, ValueSource::Unknown);

        // FIXME: There error positions counted from end of tag name - need global position
        let slice = self.start.attributes_raw();
        let decoder = self.de.reader.decoder();
        let has_value_field = self.de.has_value_field;

        if let Some(a) = self.iter.next(slice).transpose()? {
            // try getting map from attributes (key= "value")
            let (key, value) = a.into();
            self.source = ValueSource::Attribute(value.unwrap_or_default());
            seed.deserialize(EscapedDeserializer::new(
                Cow::Borrowed(&slice[key]),
                decoder,
                false,
            ))
            .map(Some)
        } else {
            // try getting from events (<key>value</key>)
            match self.de.peek()? {
                DeEvent::Text(_) | DeEvent::CData(_) => {
                    self.source = ValueSource::Text;
                    // Deserialize `key` from special attribute name which means
                    // that value should be taken from the text content of the
                    // XML node
                    seed.deserialize(INNER_VALUE.into_deserializer()).map(Some)
                }
                // Used to deserialize collections of enums, like:
                // <root>
                //
                //
                //   <C/>
                // </root>
                //
                // into
                //
                // enum Enum { A, B, С }
                // struct Root {
                //     #[serde(rename = "$value")]
                //     items: Vec<Enum>,
                // }
                // TODO: This should be handled by #[serde(flatten)]
                // See https://github.com/serde-rs/serde/issues/1905
                DeEvent::Start(_) if has_value_field => {
                    self.source = ValueSource::Text;
                    seed.deserialize(INNER_VALUE.into_deserializer()).map(Some)
                }
                DeEvent::Start(e) => {
                    self.source = ValueSource::Nested;
                    let key = if let Some(p) = self
                        .unflatten_fields
                        .iter()
                        .position(|f| e.name() == &f[UNFLATTEN_PREFIX.len()..])
                    {
                        // Used to deserialize elements, like:
                        // <root>
                        //   <xxx>test</xxx>
                        // </root>
                        //
                        // into
                        //
                        // struct Root {
                        //     #[serde(rename = "$unflatten=xxx")]
                        //     xxx: String,
                        // }
                        seed.deserialize(self.unflatten_fields.remove(p).into_deserializer())
                    } else {
                        let name = Cow::Borrowed(e.local_name());
                        seed.deserialize(EscapedDeserializer::new(name, decoder, false))
                    };
                    key.map(Some)
                }
                _ => Ok(None),
            }
        }
    }

    fn next_value_seed<K: DeserializeSeed<'de>>(
        &mut self,
        seed: K,
    ) -> Result<K::Value, Self::Error> {
        match std::mem::replace(&mut self.source, ValueSource::Unknown) {
            ValueSource::Attribute(value) => {
                let slice = self.start.attributes_raw();
                let decoder = self.de.reader.decoder();

                seed.deserialize(EscapedDeserializer::new(
                    Cow::Borrowed(&slice[value]),
                    decoder,
                    true,
                ))
            }
            ValueSource::Nested | ValueSource::Text => seed.deserialize(&mut *self.de),
            ValueSource::Unknown => Err(DeError::KeyNotRead),
        }
    }
}
<|endoftext|>


<commit_before>use super::traits::{BeaconNode, BeaconNodeError};
use protos::services::{
    BeaconBlock as GrpcBeaconBlock, ProduceBeaconBlockRequest, PublishBeaconBlockRequest,
};
use protos::services_grpc::BeaconBlockServiceClient;
use ssz::{ssz_encode, Decodable};
use types::{BeaconBlock, BeaconBlockBody, Hash256, Signature};

impl BeaconNode for BeaconBlockServiceClient {
    /// Request a Beacon Node (BN) to produce a new block at the supplied slot.
    ///
    /// Returns `None` if it is not possible to produce at the supplied slot. For example, if the
    /// BN is unable to find a parent block.
    fn produce_beacon_block(&self, slot: u64) -> Result<Option<BeaconBlock>, BeaconNodeError> {
        let mut req = ProduceBeaconBlockRequest::new();
        req.set_slot(slot);

        let reply = self
            .produce_beacon_block(&req)
            .map_err(|err| BeaconNodeError::RemoteFailure(format!("{:?}", err)))?;

        if reply.has_block() {
            let block = reply.get_block();

            let (signature, _) = Signature::ssz_decode(block.get_signature(), 0)
                .map_err(|_| BeaconNodeError::DecodeFailure)?;

            // TODO: this conversion is incomplete; fix it.
            Ok(Some(BeaconBlock {
                slot: block.get_slot(),
                parent_root: Hash256::zero(),
                state_root: Hash256::zero(),
                randao_reveal: Hash256::from(block.get_randao_reveal()),
                candidate_pow_receipt_root: Hash256::zero(),
                signature,
                body: BeaconBlockBody {
                    proposer_slashings: vec![],
                    casper_slashings: vec![],
                    attestations: vec![],
                    custody_reseeds: vec![],
                    custody_challenges: vec![],
                    custody_responses: vec![],
                    deposits: vec![],
                    exits: vec![],
                },
            }))
        } else {
            Ok(None)
        }
    }

    /// Request a Beacon Node (BN) to publish a block.
    ///
    /// Generally, this will be called after a `produce_beacon_block` call with a block that has
    /// been completed (signed) by the validator client.
    fn publish_beacon_block(&self, block: BeaconBlock) -> Result<bool, BeaconNodeError> {
        let mut req = PublishBeaconBlockRequest::new();

        // TODO: this conversion is incomplete; fix it.
        let mut grpc_block = GrpcBeaconBlock::new();
        grpc_block.set_slot(block.slot);
        grpc_block.set_block_root(vec![0]);
        grpc_block.set_randao_reveal(block.randao_reveal.to_vec());
        grpc_block.set_signature(ssz_encode(&block.signature));

        req.set_block(grpc_block);

        let reply = self
            .publish_beacon_block(&req)
            .map_err(|err| BeaconNodeError::RemoteFailure(format!("{:?}", err)))?;

        Ok(reply.get_success())
    }
}
<commit_msg>Update validator client for types change<commit_after>use super::traits::{BeaconNode, BeaconNodeError};
use protos::services::{
    BeaconBlock as GrpcBeaconBlock, ProduceBeaconBlockRequest, PublishBeaconBlockRequest,
};
use protos::services_grpc::BeaconBlockServiceClient;
use ssz::{ssz_encode, Decodable};
use types::{BeaconBlock, BeaconBlockBody, Hash256, Signature};

impl BeaconNode for BeaconBlockServiceClient {
    /// Request a Beacon Node (BN) to produce a new block at the supplied slot.
    ///
    /// Returns `None` if it is not possible to produce at the supplied slot. For example, if the
    /// BN is unable to find a parent block.
    fn produce_beacon_block(&self, slot: u64) -> Result<Option<BeaconBlock>, BeaconNodeError> {
        let mut req = ProduceBeaconBlockRequest::new();
        req.set_slot(slot);

        let reply = self
            .produce_beacon_block(&req)
            .map_err(|err| BeaconNodeError::RemoteFailure(format!("{:?}", err)))?;

        if reply.has_block() {
            let block = reply.get_block();

            let (signature, _) = Signature::ssz_decode(block.get_signature(), 0)
                .map_err(|_| BeaconNodeError::DecodeFailure)?;

            let (randao_reveal, _) = Signature::ssz_decode(block.get_randao_reveal(), 0)
                .map_err(|_| BeaconNodeError::DecodeFailure)?;

            // TODO: this conversion is incomplete; fix it.
            Ok(Some(BeaconBlock {
                slot: block.get_slot(),
                parent_root: Hash256::zero(),
                state_root: Hash256::zero(),
                randao_reveal,
                candidate_pow_receipt_root: Hash256::zero(),
                signature,
                body: BeaconBlockBody {
                    proposer_slashings: vec![],
                    casper_slashings: vec![],
                    attestations: vec![],
                    custody_reseeds: vec![],
                    custody_challenges: vec![],
                    custody_responses: vec![],
                    deposits: vec![],
                    exits: vec![],
                },
            }))
        } else {
            Ok(None)
        }
    }

    /// Request a Beacon Node (BN) to publish a block.
    ///
    /// Generally, this will be called after a `produce_beacon_block` call with a block that has
    /// been completed (signed) by the validator client.
    fn publish_beacon_block(&self, block: BeaconBlock) -> Result<bool, BeaconNodeError> {
        let mut req = PublishBeaconBlockRequest::new();

        // TODO: this conversion is incomplete; fix it.
        let mut grpc_block = GrpcBeaconBlock::new();
        grpc_block.set_slot(block.slot);
        grpc_block.set_block_root(vec![0]);
        grpc_block.set_randao_reveal(block.randao_reveal.to_vec());
        grpc_block.set_signature(ssz_encode(&block.signature));

        req.set_block(grpc_block);

        let reply = self
            .publish_beacon_block(&req)
            .map_err(|err| BeaconNodeError::RemoteFailure(format!("{:?}", err)))?;

        Ok(reply.get_success())
    }
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

//! Communicating with the running kernel using devicemapper.

#[allow(dead_code, non_camel_case_types)]
mod dm_ioctl;

use std::fs::File;
use std::io;
use std::io::{BufReader, BufRead};
use std::os::unix::io::AsRawFd;
use std::mem;
use std::slice;
use std::slice::bytes::copy_memory;
use std::io::Error;
use std::io::ErrorKind::Other;
use std::collections::BTreeSet;

use byteorder::{LittleEndian, ByteOrder};

use nix::sys::ioctl;

use dm::dm_ioctl as dmi;
use lv::LV;
use vg::VG;
use util::align_to;
use pv::Device;

const DM_IOCTL: u8 = 0xfd;
const DM_CTL_PATH: &'static str= "/dev/mapper/control";

const DM_VERSION_MAJOR: u32 = 4;
const DM_VERSION_MINOR: u32 = 30;
const DM_VERSION_PATCHLEVEL: u32 = 0;

// Status bits
//const DM_READONLY_FLAG: u32 = 1;
const DM_SUSPEND_FLAG: u32 = 2;
//const DM_PERSISTENT_DEV_FLAG: u32 = 8;

/// Context needed for communicating with devicemapper.
pub struct DM<'a> {
    file: File,
    vg: &'a VG,
}

impl <'a> DM<'a> {
    /// Create a new context for communicating about a given VG with DM.
    pub fn new(vg: &'a VG) -> io::Result<Self> {
        Ok(DM {
            file: try!(File::open(DM_CTL_PATH)),
            vg: vg,
        })
    }

    /// Determine if a major number is a major number used by DM.
    pub fn is_dm_major(major: u32) -> bool {
        let mut set = BTreeSet::new();

        let f = File::open("/proc/devices")
            .ok().expect("Could not open /proc/devices");

        let reader = BufReader::new(f);

        for line in reader.lines()
            .filter_map(|x| x.ok())
            .skip_while(|x| x != "Block devices:")
            .skip(1) {
                let spl: Vec<_> = line.split_whitespace().collect();

                if spl[1] == "device-mapper" {
                    set.insert(spl[0].parse::<u32>().unwrap());
                }
            }

        set.contains(&major)
    }

    fn initialize_hdr(hdr: &mut dmi::Struct_dm_ioctl) -> () {
        hdr.version[0] = DM_VERSION_MAJOR;
        hdr.version[1] = DM_VERSION_MINOR;
        hdr.version[2] = DM_VERSION_PATCHLEVEL;

        hdr.data_start = mem::size_of::<dmi::Struct_dm_ioctl>() as u32;
    }

    fn hdr_set_name(hdr: &mut dmi::Struct_dm_ioctl, vg_name: &str, lv_name: &str) -> () {
        let name = format!("{}-{}", vg_name.replace("-", "--"),
                           lv_name.replace("-", "--"));
        let name_dest: &mut [u8; 128] = unsafe { mem::transmute(&mut hdr.name) };
        copy_memory(name.as_bytes(), &mut name_dest[..]);
    }

    fn hdr_set_uuid(hdr: &mut dmi::Struct_dm_ioctl, vg_uuid: &str, lv_uuid: &str) -> () {
        let uuid = format!("LVM-{}{}", vg_uuid.replace("-", ""), lv_uuid.replace("-", ""));
        let uuid_dest: &mut [u8; 129] = unsafe { mem::transmute(&mut hdr.uuid) };
        copy_memory(uuid.as_bytes(), &mut uuid_dest[..]);
    }

    /// Devicemapper version information: Major, Minor, and patchlevel versions.
    pub fn get_version(&self) -> io::Result<(u32, u32, u32)> {

        let mut hdr: dmi::Struct_dm_ioctl = Default::default();
        hdr.version[0] = DM_VERSION_MAJOR;
        hdr.version[1] = DM_VERSION_MINOR;
        hdr.version[2] = DM_VERSION_PATCHLEVEL;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_VERSION_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((io::Error::last_os_error())),
            _ => {},
        };

        Ok((hdr.version[0], hdr.version[1], hdr.version[2]))
    }

    /// Returns a list of tuples containing DM device names and their major/minor
    /// device numbers.
    pub fn list_devices(&self) -> io::Result<Vec<(String, Device)>> {
        let mut buf = [0u8; 16 * 1024];
        let mut hdr: &mut dmi::Struct_dm_ioctl = unsafe {mem::transmute(&mut buf)};

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = buf.len() as u32;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_LIST_DEVICES_CMD as u8, buf.len());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut buf) } {
            Err(_) => return Err((io::Error::last_os_error())),
            _ => {},
        };

        let mut devs = Vec::new();
        if (hdr.data_size - hdr.data_start as u32) != 0 {
            let mut result = &buf[hdr.data_start as usize..];

            loop {
                let slc = slice_to_null(&result[12..]).expect("Bad data from ioctl");
                let devno = LittleEndian::read_u64(&result[..8]);
                devs.push((String::from_utf8_lossy(slc).into_owned(), Device::from(devno)));

                let next = LittleEndian::read_u32(&result[8..12]);
                if next == 0 { break }

                result = &result[next as usize..];
            }
        }

        Ok(devs)
    }

    fn create_device(&self, lv: &mut LV) -> io::Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        Self::hdr_set_name(&mut hdr, &self.vg.name, &lv.name);
        Self::hdr_set_uuid(&mut hdr, &self.vg.id, &lv.id);
        hdr.data_size = hdr.data_start;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_CREATE_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((io::Error::last_os_error())),
            _ => { }
        };

        lv.device = Some(Device::from(hdr.dev));

        Ok(())
    }

    fn remove_device(&self, lv: &LV) -> io::Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = hdr.data_start;
        Self::hdr_set_name(&mut hdr, &self.vg.name, &lv.name);

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_REMOVE_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((io::Error::last_os_error())),
            _ => Ok(())
        }
    }

    fn load_device(&self, lv: &LV) -> io::Result<()> {
        let sectors_per_extent = self.vg.extent_size;
        let mut targs = Vec::new();

        // Construct targets first, since we need to know how many & size
        // before initializing the header.
        for seg in &lv.segments {

            let seg_ty = {
                if seg.ty == "striped" && seg.stripes.len() == 1 {
                    &b"linear"[..]
                } else {
                    seg.ty.as_bytes()
                }
            };

            for &(ref pvname, ref loc) in &seg.stripes {
                let err = || Error::new(Other, "dm load_device error");
                let pv = try!(self.vg.pvs.get(pvname).ok_or(err()));

                let mut targ: dmi::Struct_dm_target_spec = Default::default();
                targ.sector_start = seg.start_extent * sectors_per_extent;
                targ.length = seg.extent_count * sectors_per_extent;
                targ.status = 0;

                let mut dst: &mut [u8] = unsafe {
                    mem::transmute(&mut targ.target_type[..])
                };
                copy_memory(seg_ty, &mut dst);

                let mut params = Vec::new();
                // TODO: only works for linear
                params.extend(
                    format!("{}:{} {}",
                            pv.device.major,
                            pv.device.minor,
                            (loc * sectors_per_extent) + pv.pe_start).as_bytes());

                let pad_bytes = align_to(
                    params.len() + 1usize, 8usize) - params.len();
                params.extend(vec![0; pad_bytes]);

                targ.next = (mem::size_of::<dmi::Struct_dm_target_spec>()
                           + params.len()) as u32;

                targs.push((targ, params));
            }
        }

        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        Self::hdr_set_name(&mut hdr, &self.vg.name, &lv.name);

        hdr.data_start = mem::size_of::<dmi::Struct_dm_ioctl>() as u32;
        hdr.data_size = hdr.data_start + targs.iter()
            .map(|&(t, _)| t.next)
            .sum::<u32>();
        hdr.target_count = targs.len() as u32;

        // Flatten into buf
        let mut buf: Vec<u8> = Vec::with_capacity(hdr.data_size as usize);
        unsafe {
            let ptr: *mut u8 = mem::transmute(&mut hdr);
            let slc = slice::from_raw_parts(ptr, hdr.data_start as usize);
            buf.extend(slc);
        }

        for (targ, param) in targs {
            unsafe {
                let ptr: *mut u8 = mem::transmute(&targ);
                let slc = slice::from_raw_parts(
                    ptr, mem::size_of::<dmi::Struct_dm_target_spec>());
                buf.extend(slc);
            }

            buf.extend(&param);
        }

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_TABLE_LOAD_CMD as u8, buf.len());

        match unsafe { ioctl::read_into_ptr(self.file.as_raw_fd(), op, buf.as_mut_ptr()) } {
            Err(_) => return Err((io::Error::last_os_error())),
            _ => Ok(())
        }
    }

    fn suspend_device(&self, lv: &LV) -> io::Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = hdr.data_start;
        Self::hdr_set_name(&mut hdr, &self.vg.name, &lv.name);
        hdr.flags = DM_SUSPEND_FLAG;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_SUSPEND_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((io::Error::last_os_error())),
            _ => Ok(())
        }
    }

    fn resume_device(&self, lv: &LV) -> io::Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = hdr.data_start;
        Self::hdr_set_name(&mut hdr, &self.vg.name, &lv.name);
        // DM_SUSPEND_FLAG not set = resume

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_SUSPEND_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((io::Error::last_os_error())),
            _ => Ok(())
        }
    }

    /// Activate a Logical Volume.
    ///
    /// Also populates the LV's device field.
    pub fn activate_device(&self, lv: &mut LV) -> io::Result<()> {
        try!(self.create_device(lv));
        try!(self.load_device(lv));
        self.resume_device(lv)
    }

    /// Deactivate a Logical Volume.
    pub fn deactivate_device(&self, lv: &mut LV) -> io::Result<()> {
        try!(self.suspend_device(lv));
        self.remove_device(lv)
    }
}

//
// Return up to the first \0, or None
//
fn slice_to_null(slc: &[u8]) -> Option<&[u8]> {
    for (i, c) in slc.iter().enumerate() {
        if *c == b'\0' { return Some(&slc[..i]) };
    }
    None
}
<commit_msg>implement dm::list_deps<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

//! Communicating with the running kernel using devicemapper.

#[allow(dead_code, non_camel_case_types)]
mod dm_ioctl;

use std::fs::File;
use std::io;
use std::io::{BufReader, BufRead};
use std::os::unix::io::AsRawFd;
use std::mem;
use std::slice;
use std::slice::bytes::copy_memory;
use std::io::Error;
use std::io::ErrorKind::Other;
use std::collections::BTreeSet;

use byteorder::{LittleEndian, ByteOrder};

use nix::sys::ioctl;

use dm::dm_ioctl as dmi;
use lv::LV;
use vg::VG;
use util::align_to;
use pv::Device;

const DM_IOCTL: u8 = 0xfd;
const DM_CTL_PATH: &'static str= "/dev/mapper/control";

const DM_VERSION_MAJOR: u32 = 4;
const DM_VERSION_MINOR: u32 = 30;
const DM_VERSION_PATCHLEVEL: u32 = 0;

// Status bits
//const DM_READONLY_FLAG: u32 = 1;
const DM_SUSPEND_FLAG: u32 = 2;
//const DM_PERSISTENT_DEV_FLAG: u32 = 8;

/// Context needed for communicating with devicemapper.
pub struct DM<'a> {
    file: File,
    vg: &'a VG,
}

impl <'a> DM<'a> {
    /// Create a new context for communicating about a given VG with DM.
    pub fn new(vg: &'a VG) -> io::Result<Self> {
        Ok(DM {
            file: try!(File::open(DM_CTL_PATH)),
            vg: vg,
        })
    }

    /// Determine if a major number is a major number used by DM.
    pub fn is_dm_major(major: u32) -> bool {
        let mut set = BTreeSet::new();

        let f = File::open("/proc/devices")
            .ok().expect("Could not open /proc/devices");

        let reader = BufReader::new(f);

        for line in reader.lines()
            .filter_map(|x| x.ok())
            .skip_while(|x| x != "Block devices:")
            .skip(1) {
                let spl: Vec<_> = line.split_whitespace().collect();

                if spl[1] == "device-mapper" {
                    set.insert(spl[0].parse::<u32>().unwrap());
                }
            }

        set.contains(&major)
    }

    fn initialize_hdr(hdr: &mut dmi::Struct_dm_ioctl) -> () {
        hdr.version[0] = DM_VERSION_MAJOR;
        hdr.version[1] = DM_VERSION_MINOR;
        hdr.version[2] = DM_VERSION_PATCHLEVEL;

        hdr.data_start = mem::size_of::<dmi::Struct_dm_ioctl>() as u32;
    }

    fn hdr_set_name(hdr: &mut dmi::Struct_dm_ioctl, vg_name: &str, lv_name: &str) -> () {
        let name = format!("{}-{}", vg_name.replace("-", "--"),
                           lv_name.replace("-", "--"));
        let name_dest: &mut [u8; 128] = unsafe { mem::transmute(&mut hdr.name) };
        copy_memory(name.as_bytes(), &mut name_dest[..]);
    }

    fn hdr_set_uuid(hdr: &mut dmi::Struct_dm_ioctl, vg_uuid: &str, lv_uuid: &str) -> () {
        let uuid = format!("LVM-{}{}", vg_uuid.replace("-", ""), lv_uuid.replace("-", ""));
        let uuid_dest: &mut [u8; 129] = unsafe { mem::transmute(&mut hdr.uuid) };
        copy_memory(uuid.as_bytes(), &mut uuid_dest[..]);
    }

    /// Devicemapper version information: Major, Minor, and patchlevel versions.
    pub fn get_version(&self) -> io::Result<(u32, u32, u32)> {

        let mut hdr: dmi::Struct_dm_ioctl = Default::default();
        hdr.version[0] = DM_VERSION_MAJOR;
        hdr.version[1] = DM_VERSION_MINOR;
        hdr.version[2] = DM_VERSION_PATCHLEVEL;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_VERSION_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((io::Error::last_os_error())),
            _ => {},
        };

        Ok((hdr.version[0], hdr.version[1], hdr.version[2]))
    }

    /// Returns a list of tuples containing DM device names and their major/minor
    /// device numbers.
    pub fn list_devices(&self) -> io::Result<Vec<(String, Device)>> {
        let mut buf = [0u8; 16 * 1024];
        let mut hdr: &mut dmi::Struct_dm_ioctl = unsafe {mem::transmute(&mut buf)};

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = buf.len() as u32;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_LIST_DEVICES_CMD as u8, buf.len());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut buf) } {
            Err(_) => return Err((io::Error::last_os_error())),
            _ => {},
        };

        let mut devs = Vec::new();
        if (hdr.data_size - hdr.data_start as u32) != 0 {
            let mut result = &buf[hdr.data_start as usize..];

            loop {
                let slc = slice_to_null(&result[12..]).expect("Bad data from ioctl");
                let devno = LittleEndian::read_u64(&result[..8]);
                devs.push((String::from_utf8_lossy(slc).into_owned(), Device::from(devno)));

                let next = LittleEndian::read_u32(&result[8..12]);
                if next == 0 { break }

                result = &result[next as usize..];
            }
        }

        Ok(devs)
    }

    pub fn list_deps(&self) -> io::Result<Vec<Device>> {
        let mut buf = [0u8; 16 * 1024];
        let mut hdr: &mut dmi::Struct_dm_ioctl = unsafe {mem::transmute(&mut buf)};

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = buf.len() as u32;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_TABLE_DEPS_CMD as u8, buf.len());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut buf) } {
            Err(_) => return Err((io::Error::last_os_error())),
            _ => {},
        };

        // TODO: Check DM_BUFFER_FULL_FLAG for:
        // DM_DEVICE_LIST_VERSIONS, DM_DEVICE_LIST, DM_DEVICE_DEPS,
        // DM_DEVICE_STATUS, DM_DEVICE_TABLE, DM_DEVICE_WAITEVENT,
        // DM_DEVICE_TARGET_MSG

        let mut devs = Vec::new();
        if (hdr.data_size - hdr.data_start as u32) != 0 {
            let result = &buf[hdr.data_start as usize..];
            let entries = LittleEndian::read_u32(&result[..4]) as usize;

            for entry in 0..entries {
                let dev = &result[(8*entry)+8..(8*entry)+16];
                devs.push(Device::from(LittleEndian::read_u64(&dev)));
            }
        }

        Ok(devs)
    }

    fn create_device(&self, lv: &mut LV) -> io::Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        Self::hdr_set_name(&mut hdr, &self.vg.name, &lv.name);
        Self::hdr_set_uuid(&mut hdr, &self.vg.id, &lv.id);
        hdr.data_size = hdr.data_start;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_CREATE_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((io::Error::last_os_error())),
            _ => { }
        };

        lv.device = Some(Device::from(hdr.dev));

        Ok(())
    }

    fn remove_device(&self, lv: &LV) -> io::Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = hdr.data_start;
        Self::hdr_set_name(&mut hdr, &self.vg.name, &lv.name);

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_REMOVE_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((io::Error::last_os_error())),
            _ => Ok(())
        }
    }

    fn load_device(&self, lv: &LV) -> io::Result<()> {
        let sectors_per_extent = self.vg.extent_size;
        let mut targs = Vec::new();

        // Construct targets first, since we need to know how many & size
        // before initializing the header.
        for seg in &lv.segments {

            let seg_ty = {
                if seg.ty == "striped" && seg.stripes.len() == 1 {
                    &b"linear"[..]
                } else {
                    seg.ty.as_bytes()
                }
            };

            for &(ref pvname, ref loc) in &seg.stripes {
                let err = || Error::new(Other, "dm load_device error");
                let pv = try!(self.vg.pvs.get(pvname).ok_or(err()));

                let mut targ: dmi::Struct_dm_target_spec = Default::default();
                targ.sector_start = seg.start_extent * sectors_per_extent;
                targ.length = seg.extent_count * sectors_per_extent;
                targ.status = 0;

                let mut dst: &mut [u8] = unsafe {
                    mem::transmute(&mut targ.target_type[..])
                };
                copy_memory(seg_ty, &mut dst);

                let mut params = Vec::new();
                // TODO: only works for linear
                params.extend(
                    format!("{}:{} {}",
                            pv.device.major,
                            pv.device.minor,
                            (loc * sectors_per_extent) + pv.pe_start).as_bytes());

                let pad_bytes = align_to(
                    params.len() + 1usize, 8usize) - params.len();
                params.extend(vec![0; pad_bytes]);

                targ.next = (mem::size_of::<dmi::Struct_dm_target_spec>()
                           + params.len()) as u32;

                targs.push((targ, params));
            }
        }

        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        Self::hdr_set_name(&mut hdr, &self.vg.name, &lv.name);

        hdr.data_start = mem::size_of::<dmi::Struct_dm_ioctl>() as u32;
        hdr.data_size = hdr.data_start + targs.iter()
            .map(|&(t, _)| t.next)
            .sum::<u32>();
        hdr.target_count = targs.len() as u32;

        // Flatten into buf
        let mut buf: Vec<u8> = Vec::with_capacity(hdr.data_size as usize);
        unsafe {
            let ptr: *mut u8 = mem::transmute(&mut hdr);
            let slc = slice::from_raw_parts(ptr, hdr.data_start as usize);
            buf.extend(slc);
        }

        for (targ, param) in targs {
            unsafe {
                let ptr: *mut u8 = mem::transmute(&targ);
                let slc = slice::from_raw_parts(
                    ptr, mem::size_of::<dmi::Struct_dm_target_spec>());
                buf.extend(slc);
            }

            buf.extend(&param);
        }

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_TABLE_LOAD_CMD as u8, buf.len());

        match unsafe { ioctl::read_into_ptr(self.file.as_raw_fd(), op, buf.as_mut_ptr()) } {
            Err(_) => return Err((io::Error::last_os_error())),
            _ => Ok(())
        }
    }

    fn suspend_device(&self, lv: &LV) -> io::Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = hdr.data_start;
        Self::hdr_set_name(&mut hdr, &self.vg.name, &lv.name);
        hdr.flags = DM_SUSPEND_FLAG;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_SUSPEND_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((io::Error::last_os_error())),
            _ => Ok(())
        }
    }

    fn resume_device(&self, lv: &LV) -> io::Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = hdr.data_start;
        Self::hdr_set_name(&mut hdr, &self.vg.name, &lv.name);
        // DM_SUSPEND_FLAG not set = resume

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_SUSPEND_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((io::Error::last_os_error())),
            _ => Ok(())
        }
    }

    /// Activate a Logical Volume.
    ///
    /// Also populates the LV's device field.
    pub fn activate_device(&self, lv: &mut LV) -> io::Result<()> {
        try!(self.create_device(lv));
        try!(self.load_device(lv));
        self.resume_device(lv)
    }

    /// Deactivate a Logical Volume.
    pub fn deactivate_device(&self, lv: &mut LV) -> io::Result<()> {
        try!(self.suspend_device(lv));
        self.remove_device(lv)
    }
}

//
// Return up to the first \0, or None
//
fn slice_to_null(slc: &[u8]) -> Option<&[u8]> {
    for (i, c) in slc.iter().enumerate() {
        if *c == b'\0' { return Some(&slc[..i]) };
    }
    None
}
<|endoftext|>


<commit_before>use super::ValidatorHarness;
use beacon_chain::BeaconChain;
pub use beacon_chain::{CheckPoint, Error as BeaconChainError};
use bls::create_proof_of_possession;
use db::{
    stores::{BeaconBlockStore, BeaconStateStore},
    MemoryDB,
};
use fork_choice::{optimised_lmd_ghost::OptimisedLMDGhost, slow_lmd_ghost::SlowLMDGhost}; // import all the algorithms
use log::debug;
use rayon::prelude::*;
use slot_clock::TestingSlotClock;
use std::collections::HashSet;
use std::fs::File;
use std::io::prelude::*;
use std::iter::FromIterator;
use std::sync::Arc;
use types::{
    BeaconBlock, ChainSpec, Deposit, DepositData, DepositInput, Eth1Data, FreeAttestation, Hash256,
    Keypair, Slot,
};

/// The beacon chain harness simulates a single beacon node with `validator_count` validators connected
/// to it. Each validator is provided a borrow to the beacon chain, where it may read
/// information and submit blocks/attestations for processing.
///
/// This test harness is useful for testing validator and internal state transition logic. It
/// is not useful for testing that multiple beacon nodes can reach consensus.
pub struct BeaconChainHarness {
    pub db: Arc<MemoryDB>,
    pub beacon_chain: Arc<BeaconChain<MemoryDB, TestingSlotClock, OptimisedLMDGhost<MemoryDB>>>,
    pub block_store: Arc<BeaconBlockStore<MemoryDB>>,
    pub state_store: Arc<BeaconStateStore<MemoryDB>>,
    pub validators: Vec<ValidatorHarness>,
    pub spec: Arc<ChainSpec>,
}

impl BeaconChainHarness {
    /// Create a new harness with:
    ///
    /// - A keypair, `BlockProducer` and `Attester` for each validator.
    /// - A new BeaconChain struct where the given validators are in the genesis.
    pub fn new(spec: ChainSpec, validator_count: usize) -> Self {
        let db = Arc::new(MemoryDB::open());
        let block_store = Arc::new(BeaconBlockStore::new(db.clone()));
        let state_store = Arc::new(BeaconStateStore::new(db.clone()));
        let genesis_time = 1_549_935_547; // 12th Feb 2018 (arbitrary value in the past).
        let slot_clock = TestingSlotClock::new(spec.genesis_slot.as_u64());
        let fork_choice = OptimisedLMDGhost::new(block_store.clone(), state_store.clone());
        let latest_eth1_data = Eth1Data {
            deposit_root: Hash256::zero(),
            block_hash: Hash256::zero(),
        };

        debug!("Generating validator keypairs...");

        let keypairs: Vec<Keypair> = (0..validator_count)
            .collect::<Vec<usize>>()
            .par_iter()
            .map(|_| Keypair::random())
            .collect();

        debug!("Creating validator deposits...");

        let initial_validator_deposits = keypairs
            .par_iter()
            .map(|keypair| Deposit {
                branch: vec![], // branch verification is not specified.
                index: 0,       // index verification is not specified.
                deposit_data: DepositData {
                    amount: 32_000_000_000, // 32 ETH (in Gwei)
                    timestamp: genesis_time - 1,
                    deposit_input: DepositInput {
                        pubkey: keypair.pk.clone(),
                        withdrawal_credentials: Hash256::zero(), // Withdrawal not possible.
                        proof_of_possession: create_proof_of_possession(&keypair),
                    },
                },
            })
            .collect();

        debug!("Creating the BeaconChain...");

        // Create the Beacon Chain
        let beacon_chain = Arc::new(
            BeaconChain::genesis(
                state_store.clone(),
                block_store.clone(),
                slot_clock,
                genesis_time,
                latest_eth1_data,
                initial_validator_deposits,
                spec.clone(),
                fork_choice,
            )
            .unwrap(),
        );

        let spec = Arc::new(spec);

        debug!("Creating validator producer and attester instances...");

        // Spawn the test validator instances.
        let validators: Vec<ValidatorHarness> = keypairs
            .iter()
            .map(|keypair| {
                ValidatorHarness::new(keypair.clone(), beacon_chain.clone(), spec.clone())
            })
            .collect();

        debug!("Created {} ValidatorHarnesss", validators.len());

        Self {
            db,
            beacon_chain,
            block_store,
            state_store,
            validators,
            spec,
        }
    }

    /// Move the `slot_clock` for the `BeaconChain` forward one slot.
    ///
    /// This is the equivalent of advancing a system clock forward one `SLOT_DURATION`.
    ///
    /// Returns the new slot.
    pub fn increment_beacon_chain_slot(&mut self) -> Slot {
        let slot = self.beacon_chain.present_slot() + 1;

        debug!("Incrementing BeaconChain slot to {}.", slot);

        self.beacon_chain.slot_clock.set_slot(slot.as_u64());
        self.beacon_chain.advance_state(slot).unwrap();
        slot
    }

    /// Gather the `FreeAttestation`s from the valiators.
    ///
    /// Note: validators will only produce attestations _once per slot_. So, if you call this twice
    /// you'll only get attestations on the first run.
    pub fn gather_free_attesations(&mut self) -> Vec<FreeAttestation> {
        let present_slot = self.beacon_chain.present_slot();

        let attesting_validators = self
            .beacon_chain
            .state
            .read()
            .get_crosslink_committees_at_slot(present_slot, false, &self.spec)
            .unwrap()
            .iter()
            .fold(vec![], |mut acc, (committee, _slot)| {
                acc.append(&mut committee.clone());
                acc
            });
        let attesting_validators: HashSet<usize> =
            HashSet::from_iter(attesting_validators.iter().cloned());

        let free_attestations: Vec<FreeAttestation> = self
            .validators
            .par_iter_mut()
            .enumerate()
            .filter_map(|(i, validator)| {
                if attesting_validators.contains(&i) {
                    // Advance the validator slot.
                    validator.set_slot(present_slot);

                    // Prompt the validator to produce an attestation (if required).
                    validator.produce_free_attestation().ok()
                } else {
                    None
                }
            })
            .collect();

        debug!(
            "Gathered {} FreeAttestations for slot {}.",
            free_attestations.len(),
            present_slot
        );

        free_attestations
    }

    /// Get the block from the proposer for the slot.
    ///
    /// Note: the validator will only produce it _once per slot_. So, if you call this twice you'll
    /// only get a block once.
    pub fn produce_block(&mut self) -> BeaconBlock {
        let present_slot = self.beacon_chain.present_slot();

        let proposer = self.beacon_chain.block_proposer(present_slot).unwrap();

        debug!(
            "Producing block from validator #{} for slot {}.",
            proposer, present_slot
        );

        // Ensure the validators slot clock is accurate.
        self.validators[proposer].set_slot(present_slot);
        self.validators[proposer].produce_block().unwrap()
    }

    /// Advances the chain with a BeaconBlock and attestations from all validators.
    ///
    /// This is the ideal scenario for the Beacon Chain, 100% honest participation from
    /// validators.
    pub fn advance_chain_with_block(&mut self) {
        self.increment_beacon_chain_slot();

        // Produce a new block.
        let block = self.produce_block();
        debug!("Submitting block for processing...");
        &mut (self.beacon_chain).process_block(block).unwrap();
        debug!("...block processed by BeaconChain.");

        debug!("Producing free attestations...");

        // Produce new attestations.
        let free_attestations = self.gather_free_attesations();

        debug!("Processing free attestations...");

        free_attestations.par_iter().for_each(|free_attestation| {
            self.beacon_chain
                .process_free_attestation(free_attestation.clone())
                .unwrap();
        });

        debug!("Free attestations processed.");
    }

    /// Dump all blocks and states from the canonical beacon chain.
    pub fn chain_dump(&self) -> Result<Vec<CheckPoint>, BeaconChainError> {
        self.beacon_chain.chain_dump()
    }

    /// Write the output of `chain_dump` to a JSON file.
    pub fn dump_to_file(&self, filename: String, chain_dump: &[CheckPoint]) {
        let json = serde_json::to_string(chain_dump).unwrap();
        let mut file = File::create(filename).unwrap();
        file.write_all(json.as_bytes())
            .expect("Failed writing dump to file.");
    }
}
<commit_msg>Fix small bug.<commit_after>use super::ValidatorHarness;
use beacon_chain::BeaconChain;
pub use beacon_chain::{CheckPoint, Error as BeaconChainError};
use bls::create_proof_of_possession;
use db::{
    stores::{BeaconBlockStore, BeaconStateStore},
    MemoryDB,
};
use fork_choice::{optimised_lmd_ghost::OptimisedLMDGhost, slow_lmd_ghost::SlowLMDGhost}; // import all the algorithms
use log::debug;
use rayon::prelude::*;
use slot_clock::TestingSlotClock;
use std::collections::HashSet;
use std::fs::File;
use std::io::prelude::*;
use std::iter::FromIterator;
use std::sync::Arc;
use types::{
    BeaconBlock, ChainSpec, Deposit, DepositData, DepositInput, Eth1Data, FreeAttestation, Hash256,
    Keypair, Slot,
};

/// The beacon chain harness simulates a single beacon node with `validator_count` validators connected
/// to it. Each validator is provided a borrow to the beacon chain, where it may read
/// information and submit blocks/attestations for processing.
///
/// This test harness is useful for testing validator and internal state transition logic. It
/// is not useful for testing that multiple beacon nodes can reach consensus.
pub struct BeaconChainHarness {
    pub db: Arc<MemoryDB>,
    pub beacon_chain: Arc<BeaconChain<MemoryDB, TestingSlotClock, OptimisedLMDGhost<MemoryDB>>>,
    pub block_store: Arc<BeaconBlockStore<MemoryDB>>,
    pub state_store: Arc<BeaconStateStore<MemoryDB>>,
    pub validators: Vec<ValidatorHarness>,
    pub spec: Arc<ChainSpec>,
}

impl BeaconChainHarness {
    /// Create a new harness with:
    ///
    /// - A keypair, `BlockProducer` and `Attester` for each validator.
    /// - A new BeaconChain struct where the given validators are in the genesis.
    pub fn new(spec: ChainSpec, validator_count: usize) -> Self {
        let db = Arc::new(MemoryDB::open());
        let block_store = Arc::new(BeaconBlockStore::new(db.clone()));
        let state_store = Arc::new(BeaconStateStore::new(db.clone()));
        let genesis_time = 1_549_935_547; // 12th Feb 2018 (arbitrary value in the past).
        let slot_clock = TestingSlotClock::new(spec.genesis_slot.as_u64());
        let fork_choice = OptimisedLMDGhost::new(block_store.clone(), state_store.clone());
        let latest_eth1_data = Eth1Data {
            deposit_root: Hash256::zero(),
            block_hash: Hash256::zero(),
        };

        debug!("Generating validator keypairs...");

        let keypairs: Vec<Keypair> = (0..validator_count)
            .collect::<Vec<usize>>()
            .par_iter()
            .map(|_| Keypair::random())
            .collect();

        debug!("Creating validator deposits...");

        let initial_validator_deposits = keypairs
            .par_iter()
            .map(|keypair| Deposit {
                branch: vec![], // branch verification is not specified.
                index: 0,       // index verification is not specified.
                deposit_data: DepositData {
                    amount: 32_000_000_000, // 32 ETH (in Gwei)
                    timestamp: genesis_time - 1,
                    deposit_input: DepositInput {
                        pubkey: keypair.pk.clone(),
                        withdrawal_credentials: Hash256::zero(), // Withdrawal not possible.
                        proof_of_possession: create_proof_of_possession(&keypair),
                    },
                },
            })
            .collect();

        debug!("Creating the BeaconChain...");

        // Create the Beacon Chain
        let beacon_chain = Arc::new(
            BeaconChain::genesis(
                state_store.clone(),
                block_store.clone(),
                slot_clock,
                genesis_time,
                latest_eth1_data,
                initial_validator_deposits,
                spec.clone(),
                fork_choice,
            )
            .unwrap(),
        );

        let spec = Arc::new(spec);

        debug!("Creating validator producer and attester instances...");

        // Spawn the test validator instances.
        let validators: Vec<ValidatorHarness> = keypairs
            .iter()
            .map(|keypair| {
                ValidatorHarness::new(keypair.clone(), beacon_chain.clone(), spec.clone())
            })
            .collect();

        debug!("Created {} ValidatorHarnesss", validators.len());

        Self {
            db,
            beacon_chain,
            block_store,
            state_store,
            validators,
            spec,
        }
    }

    /// Move the `slot_clock` for the `BeaconChain` forward one slot.
    ///
    /// This is the equivalent of advancing a system clock forward one `SLOT_DURATION`.
    ///
    /// Returns the new slot.
    pub fn increment_beacon_chain_slot(&mut self) -> Slot {
        let slot = self.beacon_chain.present_slot() + 1;

        debug!("Incrementing BeaconChain slot to {}.", slot);

        self.beacon_chain.slot_clock.set_slot(slot.as_u64());
        self.beacon_chain.advance_state(slot).unwrap();
        slot
    }

    /// Gather the `FreeAttestation`s from the valiators.
    ///
    /// Note: validators will only produce attestations _once per slot_. So, if you call this twice
    /// you'll only get attestations on the first run.
    pub fn gather_free_attesations(&mut self) -> Vec<FreeAttestation> {
        let present_slot = self.beacon_chain.present_slot();

        let attesting_validators = self
            .beacon_chain
            .state
            .read()
            .get_crosslink_committees_at_slot(present_slot, false, &self.spec)
            .unwrap()
            .iter()
            .fold(vec![], |mut acc, (committee, _slot)| {
                acc.append(&mut committee.clone());
                acc
            });
        let attesting_validators: HashSet<usize> =
            HashSet::from_iter(attesting_validators.iter().cloned());

        let free_attestations: Vec<FreeAttestation> = self
            .validators
            .par_iter_mut()
            .enumerate()
            .filter_map(|(i, validator)| {
                if attesting_validators.contains(&i) {
                    // Advance the validator slot.
                    validator.set_slot(present_slot);

                    // Prompt the validator to produce an attestation (if required).
                    validator.produce_free_attestation().ok()
                } else {
                    None
                }
            })
            .collect();

        debug!(
            "Gathered {} FreeAttestations for slot {}.",
            free_attestations.len(),
            present_slot
        );

        free_attestations
    }

    /// Get the block from the proposer for the slot.
    ///
    /// Note: the validator will only produce it _once per slot_. So, if you call this twice you'll
    /// only get a block once.
    pub fn produce_block(&mut self) -> BeaconBlock {
        let present_slot = self.beacon_chain.present_slot();

        let proposer = self.beacon_chain.block_proposer(present_slot).unwrap();

        debug!(
            "Producing block from validator #{} for slot {}.",
            proposer, present_slot
        );

        // Ensure the validators slot clock is accurate.
        self.validators[proposer].set_slot(present_slot);
        self.validators[proposer].produce_block().unwrap()
    }

    /// Advances the chain with a BeaconBlock and attestations from all validators.
    ///
    /// This is the ideal scenario for the Beacon Chain, 100% honest participation from
    /// validators.
    pub fn advance_chain_with_block(&mut self) {
        self.increment_beacon_chain_slot();

        // Produce a new block.
        let block = self.produce_block();
        debug!("Submitting block for processing...");
        self.beacon_chain.process_block(block).unwrap();
        debug!("...block processed by BeaconChain.");

        debug!("Producing free attestations...");

        // Produce new attestations.
        let free_attestations = self.gather_free_attesations();

        debug!("Processing free attestations...");

        free_attestations.par_iter().for_each(|free_attestation| {
            self.beacon_chain
                .process_free_attestation(free_attestation.clone())
                .unwrap();
        });

        debug!("Free attestations processed.");
    }

    /// Dump all blocks and states from the canonical beacon chain.
    pub fn chain_dump(&self) -> Result<Vec<CheckPoint>, BeaconChainError> {
        self.beacon_chain.chain_dump()
    }

    /// Write the output of `chain_dump` to a JSON file.
    pub fn dump_to_file(&self, filename: String, chain_dump: &[CheckPoint]) {
        let json = serde_json::to_string(chain_dump).unwrap();
        let mut file = File::create(filename).unwrap();
        file.write_all(json.as_bytes())
            .expect("Failed writing dump to file.");
    }
}
<|endoftext|>


<commit_before>#![crate_name = "docopt_macros"]
#![crate_type = "dylib"]

#![feature(plugin_registrar, quote, rustc_private)]

//! This crate defines the `docopt!` macro. It is documented in the
//! documentation of the `docopt` crate.

extern crate syntax;
extern crate rustc;
extern crate docopt;

use std::borrow::Borrow;
use std::collections::HashMap;

use rustc::plugin::Registry;
use syntax::{ast, codemap};
use syntax::diagnostic::FatalError;
use syntax::ext::base::{ExtCtxt, MacResult, MacEager, DummyResult};
use syntax::ext::build::AstBuilder;
use syntax::fold::Folder;
use syntax::owned_slice::OwnedSlice;
use syntax::parse::common::SeqSep;
use syntax::parse::parser::Parser;
use syntax::parse::token;
use syntax::print::pprust;
use syntax::ptr::P;
use syntax::util::small_vector::SmallVector;

use docopt::{Docopt, ArgvMap};
use docopt::parse::{Options, Atom, Positional, Zero, One};

type PResult<T> = Result<T, FatalError>;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("docopt", expand);
}

fn expand(cx: &mut ExtCtxt, span: codemap::Span, tts: &[ast::TokenTree])
         -> Box<MacResult+'static> {
    let parsed = match MacParser::new(cx, tts).parse() {
        Ok(parsed) => parsed,
        Err(_) => return DummyResult::any(span),
    };
    parsed.items(cx)
}

/// Parsed corresponds to the result of parsing a `docopt` macro call.
/// It can be used to write a corresponding struct.
struct Parsed {
    struct_info: StructInfo,
    doc: Docopt,
    /// Overrided type annotations for struct members. May be empty.
    /// When a type annotation for an atom doesn't exist, then one is
    /// inferred automatically. It is one of: `bool`, `usize`, `String` or
    /// `Vec<String>`.
    types: HashMap<Atom, P<ast::Ty>>,
}

impl Parsed {
    /// Returns a macro result suitable for expansion.
    /// Contains two items: one for the struct and one for the struct impls.
    fn items(&self, cx: &ExtCtxt) -> Box<MacResult+'static> {
        let mut its = vec!();
        its.push(self.struct_decl(cx));

        let struct_name = self.struct_info.name;
        let full_doc = &*self.doc.parser().full_doc;
        its.push(quote_item!(cx,
            impl $struct_name {
                #[allow(dead_code)]
                fn docopt() -> docopt::Docopt {
                    // The unwrap is justified here because this code
                    // gen only happens if the Docopt usage string is valid.
                    docopt::Docopt::new($full_doc).unwrap()
                }
            }
        ).unwrap());
        MacEager::items(SmallVector::many(its))
    }

    /// Returns an item for the struct definition.
    fn struct_decl(&self, cx: &ExtCtxt) -> P<ast::Item> {
        let name = self.struct_info.name.clone();
        let vis = if self.struct_info.public { ast::Public }
                  else { ast::Inherited };
        let def = ast::StructDef {
            fields: self.struct_fields(cx),
            ctor_id: None
        };

        let mut traits = vec!["RustcDecodable".to_string()];
        traits.extend(self.struct_info.deriving.iter().cloned());
        let attrs = vec![attribute(cx, "allow", vec!["non_snake_case"]),
                         attribute(cx, "derive", traits)];

        let st = cx.item_struct(codemap::DUMMY_SP, name.clone(), def);
        cx.item(codemap::DUMMY_SP, name, attrs, st.node.clone()).map(|mut it| {
            it.vis = vis;
            it
        })
    }

    /// Returns a list of fields for the struct definition.
    /// Handles type annotations.
    fn struct_fields(&self, cx: &ExtCtxt) -> Vec<ast::StructField> {
        let mut fields: Vec<ast::StructField> = vec!();
        for (atom, opts) in self.doc.parser().descs.iter() {
            let name = ArgvMap::key_to_struct_field(&*atom.to_string());
            let ty = match self.types.get(atom) {
                None => self.pat_type(cx, atom, opts),
                Some(ty) => ty.clone(),
            };
            fields.push(self.mk_struct_field(&*name, ty));
        }
        fields
    }

    /// Returns an inferred type for a usage pattern.
    /// This is only invoked when a type annotation is not present.
    fn pat_type(&self, cx: &ExtCtxt, atom: &Atom, opts: &Options) -> P<ast::Ty> {
        let sp = codemap::DUMMY_SP;
        match (opts.repeats, &opts.arg) {
            (false, &Zero) => {
                match atom {
                    &Positional(_) => cx.ty_ident(sp, ident("String")),
                    _ => cx.ty_ident(sp, ident("bool")),
                }
            }
            (true, &Zero) => {
                match atom {
                    &Positional(_) => ty_vec_string(cx),
                    _ => cx.ty_ident(sp, ident("usize")),
                }
            }
            (false, &One(_)) => cx.ty_ident(sp, ident("String")),
            (true, &One(_)) => ty_vec_string(cx),
        }
    }

    /// Creates a struct field from a member name and type.
    fn mk_struct_field(&self, name: &str, ty: P<ast::Ty>) -> ast::StructField {
        codemap::dummy_spanned(ast::StructField_ {
            kind: ast::NamedField(ident(name), ast::Public),
            id: ast::DUMMY_NODE_ID,
            ty: ty,
            attrs: vec!(),
        })
    }
}

/// State for parsing a `docopt` macro invocation.
struct MacParser<'a, 'b:'a> {
    cx: &'a mut ExtCtxt<'b>,
    p: Parser<'b>,
}

impl<'a, 'b> MacParser<'a, 'b> {
    fn new(cx: &'a mut ExtCtxt<'b>, tts: &[ast::TokenTree]) -> MacParser<'a, 'b> {
        let p = cx.new_parser_from_tts(tts);
        MacParser { cx: cx, p: p }
    }

    /// Main entry point for parsing arguments to `docopt` macro.
    /// First looks for an identifier for the struct name.
    /// Second, a string containing the docopt usage patterns.
    /// Third, an optional list of type annotations.
    fn parse(&mut self) -> PResult<Parsed> {
        if self.p.token == token::Eof {
            self.cx.span_err(self.cx.call_site(), "macro expects arguments");
            return Err(FatalError);
        }
        let struct_info = try!(self.parse_struct_info());
        let docstr = try!(self.parse_str());

        let sep = SeqSep {
            sep: Some(token::Comma),
            trailing_sep_allowed: true,
        };
        let types = try!(self.p.parse_seq_to_end(
            &token::Eof, sep, |p| MacParser::parse_type_annotation(p)
        )).into_iter()
          .map(|(ident, ty)| {
              let field_name = ident.to_string();
              let key = ArgvMap::struct_field_to_key(&*field_name);
              (Atom::new(&*key), ty)
           })
          .collect::<HashMap<Atom, P<ast::Ty>>>();
        try!(self.p.expect(&token::Eof));

        // This config does not matter because we're only asking for the
        // usage patterns in the Docopt string. The configuration does not
        // affect the retrieval of usage patterns.
        let doc = match Docopt::new(docstr) {
            Ok(doc) => doc,
            Err(err) => {
                self.cx.span_err(self.cx.call_site(),
                                 &*format!("Invalid Docopt usage: {}", err));
                return Err(FatalError);
            }
        };
        Ok(Parsed {
            struct_info: struct_info,
            doc: doc,
            types: types,
        })
    }

    /// Parses a single string literal. On failure, an error is logged and
    /// unit is returned.
    fn parse_str(&mut self) -> PResult<String> {
        fn lit_is_str(lit: &ast::Lit) -> bool {
            match lit.node {
                ast::LitStr(_, _) => true,
                _ => false,
            }
        }
        fn lit_to_string(lit: &ast::Lit) -> String {
            match lit.node {
                ast::LitStr(ref s, _) => s.to_string(),
                _ => panic!("BUG: expected string literal"),
            }
        }
        let exp = self.cx.expander().fold_expr(self.p.parse_expr());
        let s = match exp.node {
            ast::ExprLit(ref lit) if lit_is_str(&**lit) => {
                lit_to_string(&**lit)
            }
            _ => {
                let err = format!("Expected string literal but got {}",
                                  pprust::expr_to_string(&*exp));
                self.cx.span_err(exp.span, &*err);
                return Err(FatalError);
            }
        };
        try!(self.p.bump());
        Ok(s)
    }

    /// Parses a type annotation in a `docopt` invocation of the form
    /// `ident: Ty`.
    /// Note that this is a static method as it is used as a HOF.
    fn parse_type_annotation(p: &mut Parser)
                            -> PResult<(ast::Ident, P<ast::Ty>)> {
        let ident = try!(p.parse_ident());
        try!(p.expect(&token::Colon));
        let ty = p.parse_ty();
        Ok((ident, ty))
    }

    /// Parses struct information, like visibility, name and deriving.
    fn parse_struct_info(&mut self) -> PResult<StructInfo> {
        let public = try!(self.p.eat_keyword(token::keywords::Pub));
        let mut info = StructInfo {
            name: try!(self.p.parse_ident()),
            public: public,
            deriving: vec![],
        };
        if try!(self.p.eat(&token::Comma)) { return Ok(info); }
        let deriving = try!(self.p.parse_ident());
        if deriving.name.as_str() != "derive" {
            let err = format!("Expected 'derive' keyword but got '{}'",
                              deriving);
            self.cx.span_err(self.cx.call_site(), &*err);
            return Err(FatalError);
        }
        while !try!(self.p.eat(&token::Comma)) {
            info.deriving.push(
                try!(self.p.parse_ident()).name.to_string());
        }
        Ok(info)
    }
}

struct StructInfo {
    name: ast::Ident,
    public: bool,
    deriving: Vec<String>,
}

// Convenience functions for building intermediate values.

fn ident(s: &str) -> ast::Ident {
    ast::Ident::with_empty_ctxt(token::intern(s))
}

fn attribute<S, T>(cx: &ExtCtxt, name: S, items: Vec<T>) -> ast::Attribute
            where S: Borrow<str>, T: Borrow<str> {
    let sp = codemap::DUMMY_SP;
    let its = items.into_iter().map(|s| meta_item(cx, s.borrow())).collect();
    let mi = cx.meta_list(sp, intern(name.borrow()), its);
    cx.attribute(sp, mi)
}

fn meta_item(cx: &ExtCtxt, s: &str) -> P<ast::MetaItem> {
    cx.meta_word(codemap::DUMMY_SP, intern(s))
}

fn intern(s: &str) -> token::InternedString {
    token::intern_and_get_ident(s)
}

fn ty_vec_string(cx: &ExtCtxt) -> P<ast::Ty> {
    let sp = codemap::DUMMY_SP;
    let tystr = ast::AngleBracketedParameterData {
        lifetimes: vec![],
        types: OwnedSlice::from_vec(vec![cx.ty_ident(sp, ident("String"))]),
        bindings: OwnedSlice::empty(),
    };
    cx.ty_path(ast::Path {
        span: sp,
        global: false,
        segments: vec![ast::PathSegment {
            identifier: ident("Vec"),
            parameters: ast::PathParameters::AngleBracketedParameters(tystr),
        }]
    })
}
<commit_msg>Fixes #143.<commit_after>#![crate_name = "docopt_macros"]
#![crate_type = "dylib"]

#![feature(plugin_registrar, quote, rustc_private)]

//! This crate defines the `docopt!` macro. It is documented in the
//! documentation of the `docopt` crate.

extern crate syntax;
extern crate rustc;
extern crate docopt;

use std::borrow::Borrow;
use std::collections::HashMap;

use rustc::plugin::Registry;
use syntax::{ast, codemap};
use syntax::diagnostic::FatalError;
use syntax::ext::base::{ExtCtxt, MacResult, MacEager, DummyResult};
use syntax::ext::build::AstBuilder;
use syntax::fold::Folder;
use syntax::owned_slice::OwnedSlice;
use syntax::parse::common::SeqSep;
use syntax::parse::parser::Parser;
use syntax::parse::token;
use syntax::print::pprust;
use syntax::ptr::P;
use syntax::util::small_vector::SmallVector;

use docopt::{Docopt, ArgvMap};
use docopt::parse::{Options, Atom, Positional, Zero, One};

type PResult<T> = Result<T, FatalError>;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("docopt", expand);
}

fn expand(cx: &mut ExtCtxt, span: codemap::Span, tts: &[ast::TokenTree])
         -> Box<MacResult+'static> {
    let parsed = match MacParser::new(cx, tts).parse() {
        Ok(parsed) => parsed,
        Err(_) => return DummyResult::any(span),
    };
    parsed.items(cx)
}

/// Parsed corresponds to the result of parsing a `docopt` macro call.
/// It can be used to write a corresponding struct.
struct Parsed {
    struct_info: StructInfo,
    doc: Docopt,
    /// Overrided type annotations for struct members. May be empty.
    /// When a type annotation for an atom doesn't exist, then one is
    /// inferred automatically. It is one of: `bool`, `usize`, `String` or
    /// `Vec<String>`.
    types: HashMap<Atom, P<ast::Ty>>,
}

impl Parsed {
    /// Returns a macro result suitable for expansion.
    /// Contains two items: one for the struct and one for the struct impls.
    fn items(&self, cx: &ExtCtxt) -> Box<MacResult+'static> {
        let mut its = vec!();
        its.push(self.struct_decl(cx));

        let struct_name = self.struct_info.name;
        let full_doc = &*self.doc.parser().full_doc;
        its.push(quote_item!(cx,
            impl $struct_name {
                #[allow(dead_code)]
                fn docopt() -> docopt::Docopt {
                    // The unwrap is justified here because this code
                    // gen only happens if the Docopt usage string is valid.
                    docopt::Docopt::new($full_doc).unwrap()
                }
            }
        ).unwrap());
        MacEager::items(SmallVector::many(its))
    }

    /// Returns an item for the struct definition.
    fn struct_decl(&self, cx: &ExtCtxt) -> P<ast::Item> {
        let name = self.struct_info.name.clone();
        let vis = if self.struct_info.public { ast::Public }
                  else { ast::Inherited };
        let def = ast::VariantData::Struct(
            self.struct_fields(cx), ast::DUMMY_NODE_ID);

        let mut traits = vec!["RustcDecodable".to_string()];
        traits.extend(self.struct_info.deriving.iter().cloned());
        let attrs = vec![attribute(cx, "allow", vec!["non_snake_case"]),
                         attribute(cx, "derive", traits)];

        let st = cx.item_struct(codemap::DUMMY_SP, name.clone(), def);
        cx.item(codemap::DUMMY_SP, name, attrs, st.node.clone()).map(|mut it| {
            it.vis = vis;
            it
        })
    }

    /// Returns a list of fields for the struct definition.
    /// Handles type annotations.
    fn struct_fields(&self, cx: &ExtCtxt) -> Vec<ast::StructField> {
        let mut fields: Vec<ast::StructField> = vec!();
        for (atom, opts) in self.doc.parser().descs.iter() {
            let name = ArgvMap::key_to_struct_field(&*atom.to_string());
            let ty = match self.types.get(atom) {
                None => self.pat_type(cx, atom, opts),
                Some(ty) => ty.clone(),
            };
            fields.push(self.mk_struct_field(&*name, ty));
        }
        fields
    }

    /// Returns an inferred type for a usage pattern.
    /// This is only invoked when a type annotation is not present.
    fn pat_type(&self, cx: &ExtCtxt, atom: &Atom, opts: &Options) -> P<ast::Ty> {
        let sp = codemap::DUMMY_SP;
        match (opts.repeats, &opts.arg) {
            (false, &Zero) => {
                match atom {
                    &Positional(_) => cx.ty_ident(sp, ident("String")),
                    _ => cx.ty_ident(sp, ident("bool")),
                }
            }
            (true, &Zero) => {
                match atom {
                    &Positional(_) => ty_vec_string(cx),
                    _ => cx.ty_ident(sp, ident("usize")),
                }
            }
            (false, &One(_)) => cx.ty_ident(sp, ident("String")),
            (true, &One(_)) => ty_vec_string(cx),
        }
    }

    /// Creates a struct field from a member name and type.
    fn mk_struct_field(&self, name: &str, ty: P<ast::Ty>) -> ast::StructField {
        codemap::dummy_spanned(ast::StructField_ {
            kind: ast::NamedField(ident(name), ast::Public),
            id: ast::DUMMY_NODE_ID,
            ty: ty,
            attrs: vec!(),
        })
    }
}

/// State for parsing a `docopt` macro invocation.
struct MacParser<'a, 'b:'a> {
    cx: &'a mut ExtCtxt<'b>,
    p: Parser<'b>,
}

impl<'a, 'b> MacParser<'a, 'b> {
    fn new(cx: &'a mut ExtCtxt<'b>, tts: &[ast::TokenTree]) -> MacParser<'a, 'b> {
        let p = cx.new_parser_from_tts(tts);
        MacParser { cx: cx, p: p }
    }

    /// Main entry point for parsing arguments to `docopt` macro.
    /// First looks for an identifier for the struct name.
    /// Second, a string containing the docopt usage patterns.
    /// Third, an optional list of type annotations.
    fn parse(&mut self) -> PResult<Parsed> {
        if self.p.token == token::Eof {
            self.cx.span_err(self.cx.call_site(), "macro expects arguments");
            return Err(FatalError);
        }
        let struct_info = try!(self.parse_struct_info());
        let docstr = try!(self.parse_str());

        let sep = SeqSep {
            sep: Some(token::Comma),
            trailing_sep_allowed: true,
        };
        let types = try!(self.p.parse_seq_to_end(
            &token::Eof, sep, |p| MacParser::parse_type_annotation(p)
        )).into_iter()
          .map(|(ident, ty)| {
              let field_name = ident.to_string();
              let key = ArgvMap::struct_field_to_key(&*field_name);
              (Atom::new(&*key), ty)
           })
          .collect::<HashMap<Atom, P<ast::Ty>>>();
        try!(self.p.expect(&token::Eof));

        // This config does not matter because we're only asking for the
        // usage patterns in the Docopt string. The configuration does not
        // affect the retrieval of usage patterns.
        let doc = match Docopt::new(docstr) {
            Ok(doc) => doc,
            Err(err) => {
                self.cx.span_err(self.cx.call_site(),
                                 &*format!("Invalid Docopt usage: {}", err));
                return Err(FatalError);
            }
        };
        Ok(Parsed {
            struct_info: struct_info,
            doc: doc,
            types: types,
        })
    }

    /// Parses a single string literal. On failure, an error is logged and
    /// unit is returned.
    fn parse_str(&mut self) -> PResult<String> {
        fn lit_is_str(lit: &ast::Lit) -> bool {
            match lit.node {
                ast::LitStr(_, _) => true,
                _ => false,
            }
        }
        fn lit_to_string(lit: &ast::Lit) -> String {
            match lit.node {
                ast::LitStr(ref s, _) => s.to_string(),
                _ => panic!("BUG: expected string literal"),
            }
        }
        let exp = self.cx.expander().fold_expr(self.p.parse_expr());
        let s = match exp.node {
            ast::ExprLit(ref lit) if lit_is_str(&**lit) => {
                lit_to_string(&**lit)
            }
            _ => {
                let err = format!("Expected string literal but got {}",
                                  pprust::expr_to_string(&*exp));
                self.cx.span_err(exp.span, &*err);
                return Err(FatalError);
            }
        };
        try!(self.p.bump());
        Ok(s)
    }

    /// Parses a type annotation in a `docopt` invocation of the form
    /// `ident: Ty`.
    /// Note that this is a static method as it is used as a HOF.
    fn parse_type_annotation(p: &mut Parser)
                            -> PResult<(ast::Ident, P<ast::Ty>)> {
        let ident = try!(p.parse_ident());
        try!(p.expect(&token::Colon));
        let ty = p.parse_ty();
        Ok((ident, ty))
    }

    /// Parses struct information, like visibility, name and deriving.
    fn parse_struct_info(&mut self) -> PResult<StructInfo> {
        let public = try!(self.p.eat_keyword(token::keywords::Pub));
        let mut info = StructInfo {
            name: try!(self.p.parse_ident()),
            public: public,
            deriving: vec![],
        };
        if try!(self.p.eat(&token::Comma)) { return Ok(info); }
        let deriving = try!(self.p.parse_ident());
        if deriving.name.as_str() != "derive" {
            let err = format!("Expected 'derive' keyword but got '{}'",
                              deriving);
            self.cx.span_err(self.cx.call_site(), &*err);
            return Err(FatalError);
        }
        while !try!(self.p.eat(&token::Comma)) {
            info.deriving.push(
                try!(self.p.parse_ident()).name.to_string());
        }
        Ok(info)
    }
}

struct StructInfo {
    name: ast::Ident,
    public: bool,
    deriving: Vec<String>,
}

// Convenience functions for building intermediate values.

fn ident(s: &str) -> ast::Ident {
    ast::Ident::with_empty_ctxt(token::intern(s))
}

fn attribute<S, T>(cx: &ExtCtxt, name: S, items: Vec<T>) -> ast::Attribute
            where S: Borrow<str>, T: Borrow<str> {
    let sp = codemap::DUMMY_SP;
    let its = items.into_iter().map(|s| meta_item(cx, s.borrow())).collect();
    let mi = cx.meta_list(sp, intern(name.borrow()), its);
    cx.attribute(sp, mi)
}

fn meta_item(cx: &ExtCtxt, s: &str) -> P<ast::MetaItem> {
    cx.meta_word(codemap::DUMMY_SP, intern(s))
}

fn intern(s: &str) -> token::InternedString {
    token::intern_and_get_ident(s)
}

fn ty_vec_string(cx: &ExtCtxt) -> P<ast::Ty> {
    let sp = codemap::DUMMY_SP;
    let tystr = ast::AngleBracketedParameterData {
        lifetimes: vec![],
        types: OwnedSlice::from_vec(vec![cx.ty_ident(sp, ident("String"))]),
        bindings: OwnedSlice::empty(),
    };
    cx.ty_path(ast::Path {
        span: sp,
        global: false,
        segments: vec![ast::PathSegment {
            identifier: ident("Vec"),
            parameters: ast::PathParameters::AngleBracketedParameters(tystr),
        }]
    })
}
<|endoftext|>