<commit_before>use std::collections::BTreeMap;
use std::{env, fmt, str};
use std::error::Error;
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use toml;

pub type Dependency = (String, toml::Value);

#[derive(Debug)]
/// Catch-all error for misconfigured crates.
pub struct ManifestError;

impl Error for ManifestError {
    fn description(&self) -> &str {
        "Your Cargo.toml is either missing or incorrectly structured."
    }
}

impl fmt::Display for ManifestError {
    fn fmt(&self, format: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        format.write_str(self.description())
    }
}

enum CargoFile {
    Config,
    Lock,
}

/// A Cargo Manifest
#[derive(Debug, PartialEq)]
pub struct Manifest {
    /// Manifest contents as TOML data
    pub data: toml::Table,
}

/// If a manifest is specified, return that one, otherise perform a manifest search starting from
/// the current directory.
/// If a manifest is specified, return that one. If a path is specified, perform a manifest search
/// starting from there. If nothing is specified, start searching from the current directory
/// (`cwd`).
fn find(specified: &Option<&String>, file: CargoFile) -> Result<PathBuf, Box<Error>> {
    let file_path = specified.map(PathBuf::from);

    if let Some(path) = file_path {
        if try!(fs::metadata(&path)).is_file() {
            Ok(path)
        } else {
            search(&path, file).map_err(From::from)
        }
    } else {
        env::current_dir()
            .map_err(From::from)
            .and_then(|ref dir| search(dir, file).map_err(From::from))
    }
}

/// Search for Cargo.toml in this directory and recursively up the tree until one is found.
fn search(dir: &Path, file: CargoFile) -> Result<PathBuf, ManifestError> {
    let manifest = match file {
        CargoFile::Config => dir.join("Cargo.toml"),
        CargoFile::Lock => dir.join("Cargo.lock"),
    };

    fs::metadata(&manifest)
        .map(|_| manifest)
        .or(dir.parent().ok_or(ManifestError).and_then(|dir| search(dir, file)))
}

impl Manifest {
    /// Look for a `Cargo.toml` file
    ///
    /// Starts at the given path an goes into its parent directories until the manifest file is
    /// found. If no path is given, the process's working directory is used as a starting point.
    pub fn find_file(path: &Option<&String>) -> Result<File, Box<Error>> {
        find(path, CargoFile::Config).and_then(|path| {
            OpenOptions::new()
                .read(true)
                .write(true)
                .open(path)
                .map_err(From::from)
        })
    }

    /// Look for a `Cargo.lock` file
    ///
    /// Starts at the given path an goes into its parent directories until the manifest file is
    /// found. If no path is given, the process' working directory is used as a starting point.
    pub fn find_lock_file(path: &Option<&String>) -> Result<File, Box<Error>> {
        find(path, CargoFile::Lock).and_then(|path| {
            OpenOptions::new()
                .read(true)
                .write(true)
                .open(path)
                .map_err(From::from)
        })
    }

    /// Open the `Cargo.toml` for a path (or the process' `cwd`)
    pub fn open(path: &Option<&String>) -> Result<Manifest, Box<Error>> {
        let mut file = try!(Manifest::find_file(path));
        let mut data = String::new();
        try!(file.read_to_string(&mut data));

        data.parse()
    }

    /// Open the `Cargo.lock` for a path (or the process' `cwd`)
    pub fn open_lock_file(path: &Option<&String>) -> Result<Manifest, Box<Error>> {
        let mut file = try!(Manifest::find_lock_file(path));
        let mut data = String::new();
        try!(file.read_to_string(&mut data));

        data.parse()
    }

    /// Overwrite a file with TOML data.
    pub fn write_to_file<T: Seek + Write>(&self, file: &mut T) -> Result<(), Box<Error>> {
        try!(file.seek(SeekFrom::Start(0)));
        let mut toml = self.data.clone();

        let (proj_header, proj_data) = try!(toml.remove("package")
                                                .map(|data| ("package", data))
                                                .or_else(|| {
                                                    toml.remove("project")
                                                        .map(|data| ("project", data))
                                                })
                                                .ok_or(ManifestError));
        write!(file,
               "[{}]\n{}{}",
               proj_header,
               proj_data,
               toml::Value::Table(toml))
            .map_err(From::from)
    }

    /// Add entry to a Cargo.toml.
    fn insert_into_table(&mut self,
                         table: &str,
                         &(ref name, ref data): &Dependency)
                         -> Result<(), ManifestError> {
        let ref mut manifest = self.data;
        let entry = manifest.entry(String::from(table))
                            .or_insert(toml::Value::Table(BTreeMap::new()));
        match entry {
            &mut toml::Value::Table(ref mut deps) => {
                deps.insert(name.clone(), data.clone());
                Ok(())
            }
            _ => Err(ManifestError),
        }
    }

    /// Add entry to manifest
    pub fn add_deps(&mut self, table: &str, deps: &[Dependency]) -> Result<(), Box<Error>> {
        deps.iter()
            .map(|dep| self.insert_into_table(table, &dep))
            .collect::<Result<Vec<_>, _>>()
            .map_err(From::from)
            .map(|_| ())
    }
}

impl str::FromStr for Manifest {
    type Err = Box<Error>;

    /// Read manifest data from string
    fn from_str(input: &str) -> Result<Self, Self::Err> {
        let mut parser = toml::Parser::new(&input);

        parser.parse()
              .ok_or(parser.errors.pop())
              .map_err(Option::unwrap)
              .map_err(From::from)
              .map(|data| Manifest { data: data })
    }
}
<commit_msg>Resolve Clippy Warnings<commit_after>use std::collections::BTreeMap;
use std::{env, fmt, str};
use std::error::Error;
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use toml;

pub type Dependency = (String, toml::Value);

#[derive(Debug)]
/// Catch-all error for misconfigured crates.
pub struct ManifestError;

impl Error for ManifestError {
    fn description(&self) -> &str {
        "Your Cargo.toml is either missing or incorrectly structured."
    }
}

impl fmt::Display for ManifestError {
    fn fmt(&self, format: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        format.write_str(self.description())
    }
}

enum CargoFile {
    Config,
    Lock,
}

/// A Cargo Manifest
#[derive(Debug, PartialEq)]
pub struct Manifest {
    /// Manifest contents as TOML data
    pub data: toml::Table,
}

/// If a manifest is specified, return that one, otherise perform a manifest search starting from
/// the current directory.
/// If a manifest is specified, return that one. If a path is specified, perform a manifest search
/// starting from there. If nothing is specified, start searching from the current directory
/// (`cwd`).
fn find(specified: &Option<&String>, file: CargoFile) -> Result<PathBuf, Box<Error>> {
    let file_path = specified.map(PathBuf::from);

    if let Some(path) = file_path {
        if try!(fs::metadata(&path)).is_file() {
            Ok(path)
        } else {
            search(&path, file).map_err(From::from)
        }
    } else {
        env::current_dir()
            .map_err(From::from)
            .and_then(|ref dir| search(dir, file).map_err(From::from))
    }
}

/// Search for Cargo.toml in this directory and recursively up the tree until one is found.
fn search(dir: &Path, file: CargoFile) -> Result<PathBuf, ManifestError> {
    let manifest = match file {
        CargoFile::Config => dir.join("Cargo.toml"),
        CargoFile::Lock => dir.join("Cargo.lock"),
    };

    fs::metadata(&manifest)
        .map(|_| manifest)
        .or(dir.parent().ok_or(ManifestError).and_then(|dir| search(dir, file)))
}

impl Manifest {
    /// Look for a `Cargo.toml` file
    ///
    /// Starts at the given path an goes into its parent directories until the manifest file is
    /// found. If no path is given, the process's working directory is used as a starting point.
    pub fn find_file(path: &Option<&String>) -> Result<File, Box<Error>> {
        find(path, CargoFile::Config).and_then(|path| {
            OpenOptions::new()
                .read(true)
                .write(true)
                .open(path)
                .map_err(From::from)
        })
    }

    /// Look for a `Cargo.lock` file
    ///
    /// Starts at the given path an goes into its parent directories until the manifest file is
    /// found. If no path is given, the process' working directory is used as a starting point.
    pub fn find_lock_file(path: &Option<&String>) -> Result<File, Box<Error>> {
        find(path, CargoFile::Lock).and_then(|path| {
            OpenOptions::new()
                .read(true)
                .write(true)
                .open(path)
                .map_err(From::from)
        })
    }

    /// Open the `Cargo.toml` for a path (or the process' `cwd`)
    pub fn open(path: &Option<&String>) -> Result<Manifest, Box<Error>> {
        let mut file = try!(Manifest::find_file(path));
        let mut data = String::new();
        try!(file.read_to_string(&mut data));

        data.parse()
    }

    /// Open the `Cargo.lock` for a path (or the process' `cwd`)
    pub fn open_lock_file(path: &Option<&String>) -> Result<Manifest, Box<Error>> {
        let mut file = try!(Manifest::find_lock_file(path));
        let mut data = String::new();
        try!(file.read_to_string(&mut data));

        data.parse()
    }

    /// Overwrite a file with TOML data.
    pub fn write_to_file<T: Seek + Write>(&self, file: &mut T) -> Result<(), Box<Error>> {
        try!(file.seek(SeekFrom::Start(0)));
        let mut toml = self.data.clone();

        let (proj_header, proj_data) = try!(toml.remove("package")
                                                .map(|data| ("package", data))
                                                .or_else(|| {
                                                    toml.remove("project")
                                                        .map(|data| ("project", data))
                                                })
                                                .ok_or(ManifestError));
        write!(file,
               "[{}]\n{}{}",
               proj_header,
               proj_data,
               toml::Value::Table(toml))
            .map_err(From::from)
    }

    /// Add entry to a Cargo.toml.
    #[cfg_attr(feature = "dev", allow(toplevel_ref_arg))]
    fn insert_into_table(&mut self,
                         table: &str,
                         &(ref name, ref data): &Dependency)
                         -> Result<(), ManifestError> {
        let ref mut manifest = self.data;
        let entry = manifest.entry(String::from(table))
                            .or_insert(toml::Value::Table(BTreeMap::new()));
        match *entry {
            toml::Value::Table(ref mut deps) => {
                deps.insert(name.clone(), data.clone());
                Ok(())
            }
            _ => Err(ManifestError),
        }
    }

    /// Add entry to manifest
    pub fn add_deps(&mut self, table: &str, deps: &[Dependency]) -> Result<(), Box<Error>> {
        deps.iter()
            .map(|dep| self.insert_into_table(table, &dep))
            .collect::<Result<Vec<_>, _>>()
            .map_err(From::from)
            .map(|_| ())
    }
}

impl str::FromStr for Manifest {
    type Err = Box<Error>;

    /// Read manifest data from string
    fn from_str(input: &str) -> Result<Self, Self::Err> {
        let mut parser = toml::Parser::new(&input);

        parser.parse()
              .ok_or(parser.errors.pop())
              .map_err(Option::unwrap)
              .map_err(From::from)
              .map(|data| Manifest { data: data })
    }
}
<|endoftext|>


<commit_before>use rdkafka::consumer::{BaseConsumer, Consumer, EmptyConsumerContext};
use rdkafka::config::ClientConfig;
use rdkafka::error as rderror;
use scheduled_executor::TaskGroup;
use byteorder::{BigEndian, ReadBytesExt};

use cache::Cache;
use config::{ClusterConfig, Config};
use error::*;
use utils::read_str;

use std::collections::HashMap;
use std::error::Error;
use std::fmt;
use std::sync::{Arc, RwLock};
use std::io::Cursor;

pub type MetadataConsumer = BaseConsumer<EmptyConsumerContext>;

lazy_static! {
    pub static ref CONSUMERS: MetadataConsumerCache = MetadataConsumerCache::new();
}

pub struct MetadataConsumerCache {
    consumers: RwLock<HashMap<ClusterId, Arc<MetadataConsumer>>>,
}

impl MetadataConsumerCache {
    pub fn new() -> MetadataConsumerCache {
        MetadataConsumerCache {
            consumers: RwLock::new(HashMap::new())
        }
    }

    pub fn get(&self, cluster_id: &ClusterId) -> Option<Arc<MetadataConsumer>> {
        match self.consumers.read() {
            Ok(consumers) => (*consumers).get(cluster_id).cloned(),
            Err(_) => panic!("Poison error while reading consumer from cache")
        }
    }

    pub fn get_err(&self, cluster_id: &ClusterId) -> Result<Arc<MetadataConsumer>> {
        self.get(cluster_id).ok_or_else(|| ErrorKind::MissingConsumerError(cluster_id.clone()).into())
    }

    pub fn get_or_init(&self, cluster_id: &ClusterId, config: &ClusterConfig) -> Result<Arc<MetadataConsumer>> {
        if let Some(consumer) = self.get(cluster_id) {
            return Ok(consumer);
        }

        debug!("Creating metadata consumer for {}", cluster_id);
        let consumer = ClientConfig::new()
            .set("bootstrap.servers", &config.bootstrap_servers())
            .set("api.version.request", "true")
            .create::<MetadataConsumer>()
            .chain_err(|| format!("Consumer creation failed for {}", cluster_id))?;

        let consumer_arc = Arc::new(consumer);
        match self.consumers.write() {
            Ok(mut consumers) => (*consumers).insert(cluster_id.clone(), consumer_arc.clone()),
            Err(_) => panic!("Poison error while writing consumer to cache")
        };

        Ok(consumer_arc)
    }
}

// TODO: Use structs?
pub type BrokerId = i32;
pub type TopicName = String;

#[derive(Eq, PartialEq, Hash, Clone, Debug, PartialOrd, Ord, Serialize, Deserialize)]
pub struct ClusterId(String);

impl ClusterId {
    pub fn name(&self) -> &str {
        &self.0
    }
}

impl<'a> From<&'a str> for ClusterId {
    fn from(id: &'a str) -> ClusterId {
        ClusterId(id.to_owned())
    }
}

impl From<String> for ClusterId {
    fn from(id: String) -> ClusterId {
        ClusterId(id)
    }
}

impl fmt::Display for ClusterId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

//
// ********** METADATA **********
//

#[derive(PartialEq, Serialize, Deserialize, Debug, Clone)]
pub struct Partition {
    pub id: i32,
    pub leader: BrokerId,
    pub replicas: Vec<BrokerId>,
    pub isr: Vec<BrokerId>,
    pub error: Option<String>
}

impl Partition {
    fn new(id: i32, leader: BrokerId, mut replicas: Vec<BrokerId>, mut isr: Vec<BrokerId>, error: Option<String>) -> Partition {
        replicas.sort();
        isr.sort();
        Partition {
            id: id,
            leader: leader,
            replicas: replicas,
            isr: isr,
            error: error
        }
    }
}

#[derive(PartialEq, Serialize, Deserialize, Debug, Clone)]
pub struct Broker {
    pub id: BrokerId,
    pub hostname: String,
    pub port: i32
}

impl Broker {
    fn new(id: BrokerId, hostname: String, port: i32) -> Broker {
        Broker {
            id: id,
            hostname: hostname,
            port: port
        }
    }
}

//
// ********** GROUPS **********
//

#[derive(PartialEq, Serialize, Deserialize, Debug, Clone)]
pub struct MemberAssignment {
    pub topic: String,
    pub partitions: Vec<i32>
}

#[derive(PartialEq, Serialize, Deserialize, Debug, Clone)]
pub struct GroupMember {
    pub id: String,
    pub client_id: String,
    pub client_host: String,
    pub assignments: Vec<MemberAssignment>,
}

#[derive(PartialEq, Serialize, Deserialize, Debug, Clone)]
pub struct Group {
    pub name: String,
    pub state: String,
    pub members: Vec<GroupMember>
}

fn parse_member_assignment(payload_rdr: &mut Cursor<&[u8]>) -> Result<Vec<MemberAssignment>> {
    let _version = payload_rdr.read_i16::<BigEndian>().chain_err(|| "Failed to parse value version")?;
    let assign_len = payload_rdr.read_i32::<BigEndian>().chain_err(|| "Failed to parse assigment length")?;
    let mut assigns = Vec::with_capacity(assign_len as usize);
    for _ in 0..assign_len {
        let topic = read_str(payload_rdr).chain_err(|| "Failed to parse assignment topic name")?.to_owned();
        let partition_len = payload_rdr.read_i32::<BigEndian>().chain_err(|| "Failed to parse assignment partition length")?;
        let mut partitions = Vec::with_capacity(partition_len as usize);
        for _ in 0..partition_len {
            let partition = payload_rdr.read_i32::<BigEndian>().chain_err(|| "Failed to parse assignment partition")?;
            partitions.push(partition);
        }
        assigns.push(MemberAssignment {
            topic: topic,
            partitions: partitions
        })
    }
    Ok(assigns)
}

fn fetch_groups(consumer: &MetadataConsumer, timeout_ms: i32) -> Result<Vec<Group>> {
    let group_list = consumer.fetch_group_list(None, timeout_ms)
        .chain_err(|| "Failed to fetch consumer group list")?;

    let mut groups = Vec::new();
    for rd_group in group_list.groups() {
        let members = rd_group.members().iter()
            .map(|m| {
                let mut assigns = Vec::new();
                if rd_group.protocol_type() == "consumer" {
                    if let Some(assignment) = m.assignment() {
                        let mut payload_rdr = Cursor::new(assignment);
                        assigns = parse_member_assignment(&mut payload_rdr).expect("Parse member assignment failed");
                    }
                }
                GroupMember {
                    id: m.id().to_owned(),
                    client_id: m.client_id().to_owned(),
                    client_host: m.client_host().to_owned(),
                    assignments: assigns
                }
            })
            .collect::<Vec<_>>();
        groups.push(Group {
            name: rd_group.name().to_owned(),
            state: rd_group.state().to_owned(),
            members: members
        })
    }
    Ok(groups)
}

pub struct MetadataFetchTaskGroup {
    cache: Cache,
    config: Config,
}

impl MetadataFetchTaskGroup {
    pub fn new(cache: &Cache, config: &Config) -> MetadataFetchTaskGroup {
        MetadataFetchTaskGroup {
            cache: cache.alias(),
            config: config.clone(),
        }
    }

    fn fetch_data(&self, consumer: Arc<MetadataConsumer>, cluster_id: &ClusterId) -> Result<()> {
        let metadata = consumer.fetch_metadata(None, 120000)
            .chain_err(|| format!("Failed to fetch metadata from {}", cluster_id))?;

        // Brokers
        let mut brokers = Vec::new();
        for broker in metadata.brokers() {
            brokers.push(Broker::new(broker.id(), broker.host().to_owned(), broker.port()));
        }
        self.cache.brokers.insert(cluster_id.to_owned(), brokers)
            .chain_err(|| "Failed to insert broker information in cache")?;

        // Topics
        for topic in metadata.topics() {
            let mut partitions = Vec::with_capacity(topic.partitions().len());
            for p in topic.partitions() {
                let err_descr = p.error().map(|e| rderror::RDKafkaError::from(e).description().to_owned());
                let partition = Partition::new(p.id(), p.leader(), p.replicas().to_owned(),
                                               p.isr().to_owned(), err_descr);
                partitions.push(partition);
            }
            partitions.sort_by(|a, b| a.id.cmp(&b.id));
            // TODO: do not update if it's already there?
            self.cache.topics.insert((cluster_id.clone(), topic.name().to_owned()), partitions)
                .chain_err(|| "Failed to insert broker information in cache")?;
        }

        // Groups
        for group in fetch_groups(consumer.as_ref(), 30000)? {
            self.cache.groups.insert((cluster_id.clone(), group.name.to_owned()), group)?;
        }

        Ok(())
    }
}

impl TaskGroup for MetadataFetchTaskGroup {
    type TaskId = ClusterId;

    fn get_tasks(&self) -> Vec<ClusterId> {
        self.config.clusters.keys().cloned().collect::<Vec<_>>()
    }

    fn execute(&self, cluster_id: ClusterId) {
        match CONSUMERS.get_or_init(&cluster_id, self.config.cluster(&cluster_id).unwrap()) {
            Ok(consumer) => {
                if let Err(e) = self.fetch_data(consumer, &cluster_id) {
                    format_error_chain!(e);
                }
            },
            Err(e) => format_error_chain!(e),
        }
    }
}
<commit_msg>When deserializing assignment defaults to empty Vec<commit_after>use rdkafka::consumer::{BaseConsumer, Consumer, EmptyConsumerContext};
use rdkafka::config::ClientConfig;
use rdkafka::error as rderror;
use scheduled_executor::TaskGroup;
use byteorder::{BigEndian, ReadBytesExt};

use cache::Cache;
use config::{ClusterConfig, Config};
use error::*;
use utils::read_str;

use std::collections::HashMap;
use std::error::Error;
use std::fmt;
use std::sync::{Arc, RwLock};
use std::io::Cursor;

pub type MetadataConsumer = BaseConsumer<EmptyConsumerContext>;

lazy_static! {
    pub static ref CONSUMERS: MetadataConsumerCache = MetadataConsumerCache::new();
}

pub struct MetadataConsumerCache {
    consumers: RwLock<HashMap<ClusterId, Arc<MetadataConsumer>>>,
}

impl MetadataConsumerCache {
    pub fn new() -> MetadataConsumerCache {
        MetadataConsumerCache {
            consumers: RwLock::new(HashMap::new())
        }
    }

    pub fn get(&self, cluster_id: &ClusterId) -> Option<Arc<MetadataConsumer>> {
        match self.consumers.read() {
            Ok(consumers) => (*consumers).get(cluster_id).cloned(),
            Err(_) => panic!("Poison error while reading consumer from cache")
        }
    }

    pub fn get_err(&self, cluster_id: &ClusterId) -> Result<Arc<MetadataConsumer>> {
        self.get(cluster_id).ok_or_else(|| ErrorKind::MissingConsumerError(cluster_id.clone()).into())
    }

    pub fn get_or_init(&self, cluster_id: &ClusterId, config: &ClusterConfig) -> Result<Arc<MetadataConsumer>> {
        if let Some(consumer) = self.get(cluster_id) {
            return Ok(consumer);
        }

        debug!("Creating metadata consumer for {}", cluster_id);
        let consumer = ClientConfig::new()
            .set("bootstrap.servers", &config.bootstrap_servers())
            .set("api.version.request", "true")
            .create::<MetadataConsumer>()
            .chain_err(|| format!("Consumer creation failed for {}", cluster_id))?;

        let consumer_arc = Arc::new(consumer);
        match self.consumers.write() {
            Ok(mut consumers) => (*consumers).insert(cluster_id.clone(), consumer_arc.clone()),
            Err(_) => panic!("Poison error while writing consumer to cache")
        };

        Ok(consumer_arc)
    }
}

// TODO: Use structs?
pub type BrokerId = i32;
pub type TopicName = String;

#[derive(Eq, PartialEq, Hash, Clone, Debug, PartialOrd, Ord, Serialize, Deserialize)]
pub struct ClusterId(String);

impl ClusterId {
    pub fn name(&self) -> &str {
        &self.0
    }
}

impl<'a> From<&'a str> for ClusterId {
    fn from(id: &'a str) -> ClusterId {
        ClusterId(id.to_owned())
    }
}

impl From<String> for ClusterId {
    fn from(id: String) -> ClusterId {
        ClusterId(id)
    }
}

impl fmt::Display for ClusterId {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

//
// ********** METADATA **********
//

#[derive(PartialEq, Serialize, Deserialize, Debug, Clone)]
pub struct Partition {
    pub id: i32,
    pub leader: BrokerId,
    pub replicas: Vec<BrokerId>,
    pub isr: Vec<BrokerId>,
    pub error: Option<String>
}

impl Partition {
    fn new(id: i32, leader: BrokerId, mut replicas: Vec<BrokerId>, mut isr: Vec<BrokerId>, error: Option<String>) -> Partition {
        replicas.sort();
        isr.sort();
        Partition {
            id: id,
            leader: leader,
            replicas: replicas,
            isr: isr,
            error: error
        }
    }
}

#[derive(PartialEq, Serialize, Deserialize, Debug, Clone)]
pub struct Broker {
    pub id: BrokerId,
    pub hostname: String,
    pub port: i32
}

impl Broker {
    fn new(id: BrokerId, hostname: String, port: i32) -> Broker {
        Broker {
            id: id,
            hostname: hostname,
            port: port
        }
    }
}

//
// ********** GROUPS **********
//

#[derive(PartialEq, Serialize, Deserialize, Debug, Clone)]
pub struct MemberAssignment {
    pub topic: String,
    pub partitions: Vec<i32>
}

#[derive(PartialEq, Serialize, Deserialize, Debug, Clone)]
pub struct GroupMember {
    pub id: String,
    pub client_id: String,
    pub client_host: String,
    #[serde(default)]
    pub assignments: Vec<MemberAssignment>,
}

#[derive(PartialEq, Serialize, Deserialize, Debug, Clone)]
pub struct Group {
    pub name: String,
    pub state: String,
    pub members: Vec<GroupMember>
}

fn parse_member_assignment(payload_rdr: &mut Cursor<&[u8]>) -> Result<Vec<MemberAssignment>> {
    let _version = payload_rdr.read_i16::<BigEndian>().chain_err(|| "Failed to parse value version")?;
    let assign_len = payload_rdr.read_i32::<BigEndian>().chain_err(|| "Failed to parse assigment length")?;
    let mut assigns = Vec::with_capacity(assign_len as usize);
    for _ in 0..assign_len {
        let topic = read_str(payload_rdr).chain_err(|| "Failed to parse assignment topic name")?.to_owned();
        let partition_len = payload_rdr.read_i32::<BigEndian>().chain_err(|| "Failed to parse assignment partition length")?;
        let mut partitions = Vec::with_capacity(partition_len as usize);
        for _ in 0..partition_len {
            let partition = payload_rdr.read_i32::<BigEndian>().chain_err(|| "Failed to parse assignment partition")?;
            partitions.push(partition);
        }
        assigns.push(MemberAssignment {
            topic: topic,
            partitions: partitions
        })
    }
    Ok(assigns)
}

fn fetch_groups(consumer: &MetadataConsumer, timeout_ms: i32) -> Result<Vec<Group>> {
    let group_list = consumer.fetch_group_list(None, timeout_ms)
        .chain_err(|| "Failed to fetch consumer group list")?;

    let mut groups = Vec::new();
    for rd_group in group_list.groups() {
        let members = rd_group.members().iter()
            .map(|m| {
                let mut assigns = Vec::new();
                if rd_group.protocol_type() == "consumer" {
                    if let Some(assignment) = m.assignment() {
                        let mut payload_rdr = Cursor::new(assignment);
                        assigns = parse_member_assignment(&mut payload_rdr).expect("Parse member assignment failed");
                    }
                }
                GroupMember {
                    id: m.id().to_owned(),
                    client_id: m.client_id().to_owned(),
                    client_host: m.client_host().to_owned(),
                    assignments: assigns
                }
            })
            .collect::<Vec<_>>();
        groups.push(Group {
            name: rd_group.name().to_owned(),
            state: rd_group.state().to_owned(),
            members: members
        })
    }
    Ok(groups)
}

pub struct MetadataFetchTaskGroup {
    cache: Cache,
    config: Config,
}

impl MetadataFetchTaskGroup {
    pub fn new(cache: &Cache, config: &Config) -> MetadataFetchTaskGroup {
        MetadataFetchTaskGroup {
            cache: cache.alias(),
            config: config.clone(),
        }
    }

    fn fetch_data(&self, consumer: Arc<MetadataConsumer>, cluster_id: &ClusterId) -> Result<()> {
        let metadata = consumer.fetch_metadata(None, 120000)
            .chain_err(|| format!("Failed to fetch metadata from {}", cluster_id))?;

        // Brokers
        let mut brokers = Vec::new();
        for broker in metadata.brokers() {
            brokers.push(Broker::new(broker.id(), broker.host().to_owned(), broker.port()));
        }
        self.cache.brokers.insert(cluster_id.to_owned(), brokers)
            .chain_err(|| "Failed to insert broker information in cache")?;

        // Topics
        for topic in metadata.topics() {
            let mut partitions = Vec::with_capacity(topic.partitions().len());
            for p in topic.partitions() {
                let err_descr = p.error().map(|e| rderror::RDKafkaError::from(e).description().to_owned());
                let partition = Partition::new(p.id(), p.leader(), p.replicas().to_owned(),
                                               p.isr().to_owned(), err_descr);
                partitions.push(partition);
            }
            partitions.sort_by(|a, b| a.id.cmp(&b.id));
            // TODO: do not update if it's already there?
            self.cache.topics.insert((cluster_id.clone(), topic.name().to_owned()), partitions)
                .chain_err(|| "Failed to insert broker information in cache")?;
        }

        // Groups
        for group in fetch_groups(consumer.as_ref(), 30000)? {
            self.cache.groups.insert((cluster_id.clone(), group.name.to_owned()), group)?;
        }

        Ok(())
    }
}

impl TaskGroup for MetadataFetchTaskGroup {
    type TaskId = ClusterId;

    fn get_tasks(&self) -> Vec<ClusterId> {
        self.config.clusters.keys().cloned().collect::<Vec<_>>()
    }

    fn execute(&self, cluster_id: ClusterId) {
        match CONSUMERS.get_or_init(&cluster_id, self.config.cluster(&cluster_id).unwrap()) {
            Ok(consumer) => {
                if let Err(e) = self.fetch_data(consumer, &cluster_id) {
                    format_error_chain!(e);
                }
            },
            Err(e) => format_error_chain!(e),
        }
    }
}
<|endoftext|>


<commit_before>use scroll::{ctx::TryFromCtx, Endian};

use crate::common::*;
use crate::dbi::Module;
use crate::msf::Stream;
use crate::symbol::{InlineSiteSymbol, SymbolIndex, SymbolIter};
use crate::FallibleIterator;

mod c13;
mod constants;

#[derive(Clone, Copy, Debug)]
enum LinesSize {
    C11(usize),
    C13(usize),
}

/// This struct contains data about a single module from its module info stream.
///
/// The module info stream is where private symbols and line info is stored.
pub struct ModuleInfo<'s> {
    stream: Stream<'s>,
    symbols_size: usize,
    lines_size: LinesSize,
}

impl<'s> ModuleInfo<'s> {
    /// Parses a `ModuleInfo` from it's Module info stream data.
    pub(crate) fn parse(stream: Stream<'s>, module: &Module<'_>) -> Result<Self> {
        let info = module.info();

        let mut buf = stream.parse_buffer();
        if buf.parse_u32()? != constants::CV_SIGNATURE_C13 {
            return Err(Error::UnimplementedFeature(
                "Unsupported module info format",
            ));
        }

        let lines_size = if info.lines_size > 0 {
            LinesSize::C11(info.lines_size as usize)
        } else {
            LinesSize::C13(info.c13_lines_size as usize)
        };

        let symbols_size = info.symbols_size as usize;
        Ok(ModuleInfo {
            stream,
            symbols_size,
            lines_size,
        })
    }

    /// Get an iterator over the all symbols in this module.
    pub fn symbols(&self) -> Result<SymbolIter<'_>> {
        let mut buf = self.stream.parse_buffer();
        buf.truncate(self.symbols_size)?;
        buf.parse_u32()?;
        Ok(SymbolIter::new(buf))
    }

    /// Get an iterator over symbols starting at the given index.
    pub fn symbols_at(&self, index: SymbolIndex) -> Result<SymbolIter<'_>> {
        let mut iter = self.symbols()?;
        iter.seek(index);
        Ok(iter)
    }

    /// Returns a line program that gives access to file and line information in this module.
    pub fn line_program(&self) -> Result<LineProgram<'_>> {
        let start = self.symbols_size as usize;
        let inner = match self.lines_size {
            LinesSize::C11(_size) => return Err(Error::UnimplementedFeature("C11 line programs")),
            LinesSize::C13(size) => {
                let data = &self.stream[start..start + size];
                LineProgramInner::C13(c13::C13LineProgram::parse(data)?)
            }
        };

        Ok(LineProgram { inner })
    }
}

/// Checksum of a source file's contents.
#[derive(Clone, Debug)]
pub enum FileChecksum<'a> {
    None,
    Md5(&'a [u8]),
    Sha1(&'a [u8]),
    Sha256(&'a [u8]),
}

impl PartialEq for FileChecksum<'_> {
    fn eq(&self, other: &Self) -> bool {
        // Manual implementation to allow for None != None.
        match (self, other) {
            (&FileChecksum::Md5(lhs), &FileChecksum::Md5(rhs)) => lhs == rhs,
            (&FileChecksum::Sha1(lhs), &FileChecksum::Sha1(rhs)) => lhs == rhs,
            (&FileChecksum::Sha256(lhs), &FileChecksum::Sha256(rhs)) => lhs == rhs,
            _ => false,
        }
    }
}

/// Information record on a source file.
#[derive(Clone, Debug, PartialEq)]
pub struct FileInfo<'a> {
    /// Reference to the file name in the [string table](struct.StringTable.html).
    pub name: StringRef,

    /// Checksum of the file contents.
    pub checksum: FileChecksum<'a>,
}

/// Index of a file entry in the module.
///
/// Use the [`LineProgram`] to resolve information on the file from this offset.
#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct FileIndex(pub u32);

impl<'a> TryFromCtx<'a, Endian> for FileIndex {
    type Error = scroll::Error;
    type Size = usize;

    fn try_from_ctx(this: &'a [u8], le: Endian) -> scroll::Result<(Self, Self::Size)> {
        u32::try_from_ctx(this, le).map(|(num, s)| (Self(num), s))
    }
}

/// The kind of source construct a line info is referring to.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum LineInfoKind {
    /// A source code expression.
    Expression,
    /// A source code statement.
    Statement,
}

impl Default for LineInfoKind {
    fn default() -> Self {
        LineInfoKind::Statement
    }
}

/// Mapping of a source code offset to a source file location.
///
/// A line entry is always valid up to the subsequent entry.
#[derive(Clone, Debug, PartialEq)]
pub struct LineInfo {
    /// Source code offset.
    pub offset: PdbInternalSectionOffset,
    /// The optional length of the code.
    pub length: Option<u32>,
    /// Index of the source file in this module.
    pub file_index: FileIndex,
    /// Line number of the start of the covered range.
    pub line_start: u32,
    /// Line number of the end of the covered range.
    pub line_end: u32,
    /// Column number of the start of the covered range.
    ///
    /// This value is only present if column information is provided by the PDB. Even then, it is
    /// often zero.
    pub column_start: Option<u32>,
    /// Column number of the end of the covered range.
    ///
    /// This value is only present if column information is provided by the PDB. Even then, it is
    /// often zero.
    pub column_end: Option<u32>,
    /// Kind of this line information.
    pub kind: LineInfoKind,
}

enum LineProgramInner<'a> {
    C13(c13::C13LineProgram<'a>),
}

/// The `LineProgram` provides access to source line information for a module and its procedures.
pub struct LineProgram<'a> {
    inner: LineProgramInner<'a>,
}

impl<'a> LineProgram<'a> {
    /// Returns an iterator over all line information records of this module.
    ///
    /// Note that line records are not guaranteed to be ordered by source code offset. If a
    /// monotonic order by `PdbInternalSectionOffset` or `Rva` is required, the lines have to be
    /// sorted manually.
    pub fn lines(&self) -> LineIterator<'a> {
        match self.inner {
            LineProgramInner::C13(ref inner) => LineIterator {
                inner: LineIteratorInner::C13(inner.lines()),
            },
        }
    }

    /// Returns an iterator over all file records of this module.
    pub fn files(&self) -> FileIterator<'a> {
        match self.inner {
            LineProgramInner::C13(ref inner) => FileIterator {
                inner: FileIteratorInner::C13(inner.files()),
            },
        }
    }

    /// Returns an iterator over line records for the given section offset.
    ///
    /// This does not work with any arbitrary section offset. The iterator only returns lines if the
    /// section offset is the start of a line block. This is true for procedure or symbol offsets
    /// that specify line information. For all other offsets, the iterator will be empty.
    ///
    /// Note that line records are not guaranteed to be ordered by source code offset. If a
    /// monotonic order by `PdbInternalSectionOffset` or `Rva` is required, the lines have to be
    /// sorted manually.
    pub fn lines_at_offset(&self, offset: PdbInternalSectionOffset) -> LineIterator {
        match self.inner {
            LineProgramInner::C13(ref inner) => LineIterator {
                inner: LineIteratorInner::C13(inner.lines_at_offset(offset)),
            },
        }
    }

    /// Returns an iterator over line records for an inline site.
    pub fn inlinee_lines(
        &self,
        parent_offset: PdbInternalSectionOffset,
        inline_site: &InlineSiteSymbol<'a>,
    ) -> InlineeLineIterator<'a> {
        match self.inner {
            LineProgramInner::C13(ref inner) => InlineeLineIterator {
                inner: InlineeLineIteratorInner::C13(
                    inner.inlinee_lines(parent_offset, inline_site),
                ),
            },
        }
    }

    /// Looks up file information for the specified file.
    pub fn get_file_info(&self, offset: FileIndex) -> Result<FileInfo<'a>> {
        match self.inner {
            LineProgramInner::C13(ref inner) => inner.get_file_info(offset),
        }
    }
}

#[derive(Clone, Debug)]
enum LineIteratorInner<'a> {
    C13(c13::C13LineIterator<'a>),
}

/// An iterator over line information records in a module.
#[derive(Clone, Debug)]
pub struct LineIterator<'a> {
    inner: LineIteratorInner<'a>,
}

impl Default for LineIterator<'_> {
    fn default() -> Self {
        LineIterator {
            inner: LineIteratorInner::C13(Default::default()),
        }
    }
}

impl<'a> FallibleIterator for LineIterator<'a> {
    type Item = LineInfo;
    type Error = Error;

    fn next(&mut self) -> Result<Option<Self::Item>> {
        match self.inner {
            LineIteratorInner::C13(ref mut inner) => inner.next(),
        }
    }
}

#[derive(Clone, Debug)]
enum InlineeLineIteratorInner<'a> {
    C13(c13::C13InlineeLineIterator<'a>),
}

/// An iterator over line information records in a module.
#[derive(Clone, Debug)]
pub struct InlineeLineIterator<'a> {
    inner: InlineeLineIteratorInner<'a>,
}

impl Default for InlineeLineIterator<'_> {
    fn default() -> Self {
        InlineeLineIterator {
            inner: InlineeLineIteratorInner::C13(Default::default()),
        }
    }
}

impl<'a> FallibleIterator for InlineeLineIterator<'a> {
    type Item = LineInfo;
    type Error = Error;

    fn next(&mut self) -> Result<Option<Self::Item>> {
        match self.inner {
            InlineeLineIteratorInner::C13(ref mut inner) => inner.next(),
        }
    }
}

#[derive(Clone, Debug)]
enum FileIteratorInner<'a> {
    C13(c13::C13FileIterator<'a>),
}

/// An iterator over file records in a module.
#[derive(Clone, Debug)]
pub struct FileIterator<'a> {
    inner: FileIteratorInner<'a>,
}

impl Default for FileIterator<'_> {
    fn default() -> Self {
        FileIterator {
            inner: FileIteratorInner::C13(Default::default()),
        }
    }
}

impl<'a> FallibleIterator for FileIterator<'a> {
    type Item = FileInfo<'a>;
    type Error = Error;

    fn next(&mut self) -> Result<Option<Self::Item>> {
        match self.inner {
            FileIteratorInner::C13(ref mut inner) => inner.next(),
        }
    }
}
<commit_msg>Add more doc comments<commit_after>use scroll::{ctx::TryFromCtx, Endian};

use crate::common::*;
use crate::dbi::Module;
use crate::msf::Stream;
use crate::symbol::{InlineSiteSymbol, SymbolIndex, SymbolIter};
use crate::FallibleIterator;

mod c13;
mod constants;

#[derive(Clone, Copy, Debug)]
enum LinesSize {
    C11(usize),
    C13(usize),
}

/// This struct contains data about a single module from its module info stream.
///
/// The module info stream is where private symbols and line info is stored.
pub struct ModuleInfo<'s> {
    stream: Stream<'s>,
    symbols_size: usize,
    lines_size: LinesSize,
}

impl<'s> ModuleInfo<'s> {
    /// Parses a `ModuleInfo` from it's Module info stream data.
    pub(crate) fn parse(stream: Stream<'s>, module: &Module<'_>) -> Result<Self> {
        let info = module.info();

        let mut buf = stream.parse_buffer();
        if buf.parse_u32()? != constants::CV_SIGNATURE_C13 {
            return Err(Error::UnimplementedFeature(
                "Unsupported module info format",
            ));
        }

        let lines_size = if info.lines_size > 0 {
            LinesSize::C11(info.lines_size as usize)
        } else {
            LinesSize::C13(info.c13_lines_size as usize)
        };

        let symbols_size = info.symbols_size as usize;
        Ok(ModuleInfo {
            stream,
            symbols_size,
            lines_size,
        })
    }

    /// Get an iterator over the all symbols in this module.
    pub fn symbols(&self) -> Result<SymbolIter<'_>> {
        let mut buf = self.stream.parse_buffer();
        buf.truncate(self.symbols_size)?;
        buf.parse_u32()?;
        Ok(SymbolIter::new(buf))
    }

    /// Get an iterator over symbols starting at the given index.
    pub fn symbols_at(&self, index: SymbolIndex) -> Result<SymbolIter<'_>> {
        let mut iter = self.symbols()?;
        iter.seek(index);
        Ok(iter)
    }

    /// Returns a line program that gives access to file and line information in this module.
    pub fn line_program(&self) -> Result<LineProgram<'_>> {
        let start = self.symbols_size as usize;
        let inner = match self.lines_size {
            LinesSize::C11(_size) => return Err(Error::UnimplementedFeature("C11 line programs")),
            LinesSize::C13(size) => {
                let data = &self.stream[start..start + size];
                LineProgramInner::C13(c13::C13LineProgram::parse(data)?)
            }
        };

        Ok(LineProgram { inner })
    }
}

/// Checksum of a source file's contents.
#[derive(Clone, Debug)]
pub enum FileChecksum<'a> {
    None,
    Md5(&'a [u8]),
    Sha1(&'a [u8]),
    Sha256(&'a [u8]),
}

impl PartialEq for FileChecksum<'_> {
    fn eq(&self, other: &Self) -> bool {
        // Manual implementation to allow for None != None.
        match (self, other) {
            (&FileChecksum::Md5(lhs), &FileChecksum::Md5(rhs)) => lhs == rhs,
            (&FileChecksum::Sha1(lhs), &FileChecksum::Sha1(rhs)) => lhs == rhs,
            (&FileChecksum::Sha256(lhs), &FileChecksum::Sha256(rhs)) => lhs == rhs,
            _ => false,
        }
    }
}

/// Information record on a source file.
#[derive(Clone, Debug, PartialEq)]
pub struct FileInfo<'a> {
    /// Reference to the file name in the [string table](struct.StringTable.html).
    pub name: StringRef,

    /// Checksum of the file contents.
    pub checksum: FileChecksum<'a>,
}

/// Index of a file entry in the module.
///
/// Use the [`LineProgram`] to resolve information on the file from this offset.
#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct FileIndex(pub u32);

impl<'a> TryFromCtx<'a, Endian> for FileIndex {
    type Error = scroll::Error;
    type Size = usize;

    fn try_from_ctx(this: &'a [u8], le: Endian) -> scroll::Result<(Self, Self::Size)> {
        u32::try_from_ctx(this, le).map(|(num, s)| (Self(num), s))
    }
}

/// The kind of source construct a line info is referring to.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum LineInfoKind {
    /// A source code expression.
    Expression,
    /// A source code statement.
    Statement,
}

impl Default for LineInfoKind {
    fn default() -> Self {
        LineInfoKind::Statement
    }
}

/// Mapping of a source code offset to a source file location.
///
/// A line entry is always valid up to the subsequent entry.
#[derive(Clone, Debug, PartialEq)]
pub struct LineInfo {
    /// Source code offset.
    pub offset: PdbInternalSectionOffset,
    /// The optional length of the code.
    pub length: Option<u32>,
    /// Index of the source file in this module.
    pub file_index: FileIndex,
    /// Line number of the start of the covered range.
    pub line_start: u32,
    /// Line number of the end of the covered range.
    pub line_end: u32,
    /// Column number of the start of the covered range.
    ///
    /// This value is only present if column information is provided by the PDB. Even then, it is
    /// often zero.
    pub column_start: Option<u32>,
    /// Column number of the end of the covered range.
    ///
    /// This value is only present if column information is provided by the PDB. Even then, it is
    /// often zero.
    pub column_end: Option<u32>,
    /// Kind of this line information.
    pub kind: LineInfoKind,
}

enum LineProgramInner<'a> {
    C13(c13::C13LineProgram<'a>),
}

/// The `LineProgram` provides access to source line information for a module and its procedures.
pub struct LineProgram<'a> {
    inner: LineProgramInner<'a>,
}

impl<'a> LineProgram<'a> {
    /// Returns an iterator over all line information records of this module.
    ///
    /// Note that line records are not guaranteed to be ordered by source code offset. If a
    /// monotonic order by `PdbInternalSectionOffset` or `Rva` is required, the lines have to be
    /// sorted manually.
    pub fn lines(&self) -> LineIterator<'a> {
        match self.inner {
            LineProgramInner::C13(ref inner) => LineIterator {
                inner: LineIteratorInner::C13(inner.lines()),
            },
        }
    }

    /// Returns an iterator over all file records of this module.
    pub fn files(&self) -> FileIterator<'a> {
        match self.inner {
            LineProgramInner::C13(ref inner) => FileIterator {
                inner: FileIteratorInner::C13(inner.files()),
            },
        }
    }

    /// Returns an iterator over line records for the given section offset.
    ///
    /// This does not work with any arbitrary section offset. The iterator only returns lines if the
    /// section offset is the start of a line block. This is true for procedure or symbol offsets
    /// that specify line information. For all other offsets, the iterator will be empty.
    ///
    /// Note that line records are not guaranteed to be ordered by source code offset. If a
    /// monotonic order by `PdbInternalSectionOffset` or `Rva` is required, the lines have to be
    /// sorted manually.
    pub fn lines_at_offset(&self, offset: PdbInternalSectionOffset) -> LineIterator {
        match self.inner {
            LineProgramInner::C13(ref inner) => LineIterator {
                inner: LineIteratorInner::C13(inner.lines_at_offset(offset)),
            },
        }
    }

    /// Returns an iterator over line records for an inline site.
    ///
    /// Note that line records are not guaranteed to be ordered by source code offset. If a
    /// monotonic order by `PdbInternalSectionOffset` or `Rva` is required, the lines have to be
    /// sorted manually.
    pub fn inlinee_lines(
        &self,
        parent_offset: PdbInternalSectionOffset,
        inline_site: &InlineSiteSymbol<'a>,
    ) -> InlineeLineIterator<'a> {
        match self.inner {
            LineProgramInner::C13(ref inner) => InlineeLineIterator {
                inner: InlineeLineIteratorInner::C13(
                    inner.inlinee_lines(parent_offset, inline_site),
                ),
            },
        }
    }

    /// Looks up file information for the specified file.
    pub fn get_file_info(&self, offset: FileIndex) -> Result<FileInfo<'a>> {
        match self.inner {
            LineProgramInner::C13(ref inner) => inner.get_file_info(offset),
        }
    }
}

#[derive(Clone, Debug)]
enum LineIteratorInner<'a> {
    C13(c13::C13LineIterator<'a>),
}

/// An iterator over line information records in a module.
#[derive(Clone, Debug)]
pub struct LineIterator<'a> {
    inner: LineIteratorInner<'a>,
}

impl Default for LineIterator<'_> {
    fn default() -> Self {
        LineIterator {
            inner: LineIteratorInner::C13(Default::default()),
        }
    }
}

impl<'a> FallibleIterator for LineIterator<'a> {
    type Item = LineInfo;
    type Error = Error;

    fn next(&mut self) -> Result<Option<Self::Item>> {
        match self.inner {
            LineIteratorInner::C13(ref mut inner) => inner.next(),
        }
    }
}

#[derive(Clone, Debug)]
enum InlineeLineIteratorInner<'a> {
    C13(c13::C13InlineeLineIterator<'a>),
}

/// An iterator over line information records in a module.
#[derive(Clone, Debug)]
pub struct InlineeLineIterator<'a> {
    inner: InlineeLineIteratorInner<'a>,
}

impl Default for InlineeLineIterator<'_> {
    fn default() -> Self {
        InlineeLineIterator {
            inner: InlineeLineIteratorInner::C13(Default::default()),
        }
    }
}

impl<'a> FallibleIterator for InlineeLineIterator<'a> {
    type Item = LineInfo;
    type Error = Error;

    fn next(&mut self) -> Result<Option<Self::Item>> {
        match self.inner {
            InlineeLineIteratorInner::C13(ref mut inner) => inner.next(),
        }
    }
}

#[derive(Clone, Debug)]
enum FileIteratorInner<'a> {
    C13(c13::C13FileIterator<'a>),
}

/// An iterator over file records in a module.
#[derive(Clone, Debug)]
pub struct FileIterator<'a> {
    inner: FileIteratorInner<'a>,
}

impl Default for FileIterator<'_> {
    fn default() -> Self {
        FileIterator {
            inner: FileIteratorInner::C13(Default::default()),
        }
    }
}

impl<'a> FallibleIterator for FileIterator<'a> {
    type Item = FileInfo<'a>;
    type Error = Error;

    fn next(&mut self) -> Result<Option<Self::Item>> {
        match self.inner {
            FileIteratorInner::C13(ref mut inner) => inner.next(),
        }
    }
}
<|endoftext|>


<commit_before>use arrayvec::ArrayString;
use petgraph::graphmap::DiGraphMap;
use std::cmp::{Ord, Ordering};
use std::hash::{Hash, Hasher};

// TODO: Add support to represent use's of individual members (fn's, trairs
//       etc.).  This would require using a union type (of Mod + member?) as
//       node type.

/// Determines the maximum length of a module's path.
///
/// eg: `"my_crate::foo::bar::baz"`.
const MOD_PATH_SIZE: usize = 200;

static SEP: &'static str = "::";

/// Represents an association between modules.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Edge {
    Child,
    Dependency,
}

/// Builds a graph, `DiGraphMap<Mod, Edge>` to be specific using domain
/// specific operations.
pub struct GraphBuilder {
    graph: DiGraphMap<Mod, Edge>,
    deferred_deps: Vec<(String, String)>,
}

impl GraphBuilder {
    pub fn new() -> Self {
        Self {
            graph: DiGraphMap::new(),
            deferred_deps: vec![],
        }
    }

    /// Define a public module that has no parent.
    ///
    /// # Panics
    /// If `name` contains `"::"`.
    pub fn add_crate_root(&mut self, name: &str) {
        assert!(!name.contains(SEP));
        self.graph.add_node(Mod::new(name, Visibility::Public));
    }

    /// Define a sub-modules and associate it with its parent.
    ///
    /// # Panics
    /// - If the parent is not already defined.
    /// - If parent-child relationship for this pair is already defined.
    pub fn add_mod(&mut self, path: &str, name: &str, visibility: Visibility) {
        let parent: Mod = self.find_mod(path).unwrap();
        let node = Mod::new(&[path, SEP, name].concat(), visibility);
        self.graph.add_node(node);
        assert!(self.graph.add_edge(parent, node, Edge::Child).is_none());
    }

    pub fn add_dep(&mut self, from: &str, to: &str) {
        assert!(from != to, "Module cannot depend on itself");
        match (self.find_mod(from), self.find_mod(to)) {
            (Some(from), Some(to)) => {
                self.graph.add_edge(from, to, Edge::Dependency);
                ()
            }
            // Defer creating dependency link if
            // one of the nodes are not defined yet.
            (_, _) => self.deferred_deps.push((from.to_owned(), to.to_owned())),
        };
    }

    /// Build the graph, consuming this builder, or return an error.
    pub fn build(mut self) -> Result<DiGraphMap<Mod, Edge>, GraphError> {
        if self.deferred_deps.is_empty() {
            Ok(self.graph)
        } else {
            let (from, to) = self.deferred_deps.remove(0);
            match self.find_mod(&from) {
                Some(_) => Err(GraphError::UnknownModule(to)),
                None => Err(GraphError::UnknownModule(from)),
            }
        }
    }

    fn find_mod(&self, path: &str) -> Option<Mod> {
        self.graph.nodes().find(|m| m.path() == path)
    }
}

#[derive(Debug, PartialEq)]
pub enum GraphError {
    UnknownModule(String),
}

/// Represents a node that is a module in the graph.
#[derive(Clone, Copy, Debug)]
pub struct Mod {
    /// Because this struct needs to be `Copy`, using `String` or `&str` was
    /// out of the question.  `ArrayString` provides a sized and owned string
    /// that is backed by a byte array.
    ///
    /// See also [MOD_PATH_SIZE]
    path: ArrayString<[u8; MOD_PATH_SIZE]>,
    name_ridx: usize,
    visibility: Visibility,
}

impl Mod {
    fn new(path: &str, visibility: Visibility) -> Self {
        Self {
            path: ArrayString::<[u8; MOD_PATH_SIZE]>::from(path)
                .unwrap_or_else(|_| panic!("Module path is too long")),
            name_ridx: path.rfind(SEP).unwrap_or(0),
            visibility,
        }
    }

    fn path(&self) -> &str {
        &self.path.as_str()
    }
}

impl Eq for Mod {}

impl Hash for Mod {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.path().hash(state);
    }
}

impl Ord for Mod {
    fn cmp(&self, other: &Mod) -> Ordering {
        self.path().cmp(&other.path())
    }
}

impl PartialEq for Mod {
    fn eq(&self, other: &Mod) -> bool {
        self.path() == other.path()
    }
}

impl PartialOrd for Mod {
    fn partial_cmp(&self, other: &Mod) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Clone, Copy, Debug)]
pub enum Visibility {
    Public,
    Private,
}

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

    #[test]
    fn new_builder_produces_an_empty_directed_graph() {
        let builder = GraphBuilder::new();
        let graph: DiGraphMap<Mod, Edge> = builder.build().unwrap();
        assert_eq!(0, graph.node_count());
        assert_eq!(0, graph.edge_count());
    }

    #[test]
    fn add_crate_root_adds_a_node() {
        let mut builder = GraphBuilder::new();
        builder.add_crate_root("crate-root");
        let graph: DiGraphMap<Mod, Edge> = builder.build().unwrap();
        assert_eq!(1, graph.node_count());
        assert_eq!(0, graph.edge_count());
        assert!(graph.contains_node(Mod::new("crate-root", Visibility::Public)));
    }

    #[test]
    fn add_mod_creates_an_association_with_parent() {
        let foo: Mod = Mod::new("foo", Visibility::Public);
        let bar: Mod = Mod::new("foo::bar", Visibility::Public);
        let baz: Mod = Mod::new("foo::bar::baz", Visibility::Private);
        let mut builder = GraphBuilder::new();
        builder.add_crate_root("foo");
        builder.add_mod("foo", "bar", Visibility::Public);
        builder.add_mod("foo::bar", "baz", Visibility::Private);
        let graph: DiGraphMap<Mod, Edge> = builder.build().unwrap();
        assert_eq!(3, graph.node_count());
        assert_eq!(2, graph.edge_count());
        assert!(graph.contains_node(foo));
        assert!(graph.contains_node(bar));
        assert!(graph.contains_node(baz));
        assert_eq!(
            vec![(foo, bar, &Edge::Child), (bar, baz, &Edge::Child)],
            graph.all_edges().collect::<Vec<(Mod, Mod, &Edge)>>()
        );
    }

    #[test]
    #[should_panic(expected = "Module path is too long")]
    fn add_mod_panics_if_path_is_longer_than_limit() {
        let name = "foo";
        let path = std::str::from_utf8([114; MOD_PATH_SIZE - 4].as_ref()).unwrap(); // 'r'
        let concatenated = format!("{}::{}", path, name);
        assert!(concatenated.len() > MOD_PATH_SIZE);
        let mut builder = GraphBuilder::new();
        builder.add_crate_root(path);
        builder.add_mod(path, name, Visibility::Public);
    }

    #[test]
    #[should_panic(expected = "Module cannot depend on itself")]
    fn adding_a_dependency_to_the_same_module_panics() {
        let mut builder = GraphBuilder::new();
        builder.add_crate_root("root");
        builder.add_mod("root", "sub", Visibility::Public);
        builder.add_dep("root::sub", "root::sub");
    }

    #[test]
    fn adding_a_dependency_is_idempotent() {
        let mut builder = GraphBuilder::new();
        builder.add_crate_root("foo");
        builder.add_mod("foo", "bar", Visibility::Public);
        builder.add_mod("foo", "baz", Visibility::Private);
        builder.add_dep("foo::bar", "foo::baz");
        builder.add_dep("foo::bar", "foo::baz");
        let graph = builder.build().unwrap();
        assert_eq!(3, graph.edge_count());
    }

    #[test]
    fn add_dep_requires_both_modules_to_be_defined() {
        {
            let mut builder = GraphBuilder::new();
            builder.add_crate_root("foo");
            builder.add_mod("foo", "bar", Visibility::Public);
            builder.add_dep("foo::bar", "foo::baz");
            assert_eq!(
                Some(GraphError::UnknownModule(String::from("foo::baz"))),
                builder.build().err()
            );
        }
        {
            let mut builder = GraphBuilder::new();
            builder.add_crate_root("foo");
            builder.add_mod("foo", "baz", Visibility::Private);
            builder.add_dep("foo::bar", "foo::baz");
            assert_eq!(
                Some(GraphError::UnknownModule(String::from("foo::bar"))),
                builder.build().err()
            );
        }
        {
            let mut builder = GraphBuilder::new();
            builder.add_crate_root("foo");
            builder.add_dep("foo::bar", "foo::baz");
            assert_eq!(
                Some(GraphError::UnknownModule(String::from("foo::bar"))),
                builder.build().err()
            );
        }
    }

    // TODO: Add test for deferring dependency
}
<commit_msg>Allow adding dependencies for modules that aren't visited yet<commit_after>use arrayvec::ArrayString;
use petgraph::graphmap::DiGraphMap;
use std::cmp::{Ord, Ordering};
use std::hash::{Hash, Hasher};

// TODO: Add support to represent use's of individual members (fn's, trairs
//       etc.).  This would require using a union type (of Mod + member?) as
//       node type.

/// Determines the maximum length of a module's path.
///
/// eg: `"my_crate::foo::bar::baz"`.
const MOD_PATH_SIZE: usize = 200;

static SEP: &'static str = "::";

/// Represents an association between modules.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Edge {
    Child,
    Dependency,
}

/// Builds a graph, `DiGraphMap<Mod, Edge>` to be specific using domain
/// specific operations.
pub struct GraphBuilder {
    graph: DiGraphMap<Mod, Edge>,
    deferred_deps: Vec<(String, String)>,
}

impl GraphBuilder {
    pub fn new() -> Self {
        Self {
            graph: DiGraphMap::new(),
            deferred_deps: vec![],
        }
    }

    /// Define a public module that has no parent.
    ///
    /// # Panics
    /// If `name` contains `"::"`.
    pub fn add_crate_root(&mut self, name: &str) {
        assert!(!name.contains(SEP));
        self.graph.add_node(Mod::new(name, Visibility::Public));
    }

    /// Define a sub-modules and associate it with its parent.
    ///
    /// # Panics
    /// - If the parent is not already defined.
    /// - If parent-child relationship for this pair is already defined.
    pub fn add_mod(&mut self, path: &str, name: &str, visibility: Visibility) {
        let parent: Mod = self.find_mod(path).unwrap();
        let node = Mod::new(&[path, SEP, name].concat(), visibility);
        self.graph.add_node(node);
        assert!(self.graph.add_edge(parent, node, Edge::Child).is_none());
        self.apply_deferred();
    }

    pub fn add_dep(&mut self, from: &str, to: &str) {
        assert!(from != to, "Module cannot depend on itself");
        match (self.find_mod(from), self.find_mod(to)) {
            (Some(from), Some(to)) => {
                self.graph.add_edge(from, to, Edge::Dependency);
                ()
            }
            (None, _) => panic!("Trying to add dependency from an unknown module"),
            // Defer creating dependency link if
            // one of the nodes are not defined yet.
            (_, _) => self.deferred_deps.push((from.to_owned(), to.to_owned())),
        };
    }

    /// Build the graph, consuming this builder, or return an error.
    pub fn build(mut self) -> Result<DiGraphMap<Mod, Edge>, GraphError> {
        if self.deferred_deps.is_empty() {
            Ok(self.graph)
        } else {
            let (from, to) = self.deferred_deps.remove(0);
            match self.find_mod(&from) {
                Some(_) => Err(GraphError::UnknownModule(to)),
                None => Err(GraphError::UnknownModule(from)),
            }
        }
    }

    fn apply_deferred(&mut self) {
        let deferred: Vec<(String, String)> = self.deferred_deps.drain(..).collect();
        for (from, to) in deferred {
            // We are only checking `to` because a dependent module
            // needs to be defined before the dependency is defined.
            if self.find_mod(&to).is_some() {
                self.add_dep(&from, &to);
            } else {
                self.deferred_deps.push((from, to));
            }
        }
    }

    fn find_mod(&self, path: &str) -> Option<Mod> {
        self.graph.nodes().find(|m| m.path() == path)
    }
}

#[derive(Debug, PartialEq)]
pub enum GraphError {
    UnknownModule(String),
}

/// Represents a node that is a module in the graph.
#[derive(Clone, Copy, Debug)]
pub struct Mod {
    /// Because this struct needs to be `Copy`, using `String` or `&str` was
    /// out of the question.  `ArrayString` provides a sized and owned string
    /// that is backed by a byte array.
    ///
    /// See also [MOD_PATH_SIZE]
    path: ArrayString<[u8; MOD_PATH_SIZE]>,
    name_ridx: usize,
    visibility: Visibility,
}

impl Mod {
    fn new(path: &str, visibility: Visibility) -> Self {
        Self {
            path: ArrayString::<[u8; MOD_PATH_SIZE]>::from(path)
                .unwrap_or_else(|_| panic!("Module path is too long")),
            name_ridx: path.rfind(SEP).unwrap_or(0),
            visibility,
        }
    }

    fn path(&self) -> &str {
        &self.path.as_str()
    }
}

impl Eq for Mod {}

impl Hash for Mod {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.path().hash(state);
    }
}

impl Ord for Mod {
    fn cmp(&self, other: &Mod) -> Ordering {
        self.path().cmp(&other.path())
    }
}

impl PartialEq for Mod {
    fn eq(&self, other: &Mod) -> bool {
        self.path() == other.path()
    }
}

impl PartialOrd for Mod {
    fn partial_cmp(&self, other: &Mod) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Clone, Copy, Debug)]
pub enum Visibility {
    Public,
    Private,
}

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

    #[test]
    fn new_builder_produces_an_empty_directed_graph() {
        let builder = GraphBuilder::new();
        let graph: DiGraphMap<Mod, Edge> = builder.build().unwrap();
        assert_eq!(0, graph.node_count());
        assert_eq!(0, graph.edge_count());
    }

    #[test]
    fn add_crate_root_adds_a_node() {
        let mut builder = GraphBuilder::new();
        builder.add_crate_root("crate-root");
        let graph: DiGraphMap<Mod, Edge> = builder.build().unwrap();
        assert_eq!(1, graph.node_count());
        assert_eq!(0, graph.edge_count());
        assert!(graph.contains_node(Mod::new("crate-root", Visibility::Public)));
    }

    #[test]
    fn add_mod_creates_an_association_with_parent() {
        let foo: Mod = Mod::new("foo", Visibility::Public);
        let bar: Mod = Mod::new("foo::bar", Visibility::Public);
        let baz: Mod = Mod::new("foo::bar::baz", Visibility::Private);
        let mut builder = GraphBuilder::new();
        builder.add_crate_root("foo");
        builder.add_mod("foo", "bar", Visibility::Public);
        builder.add_mod("foo::bar", "baz", Visibility::Private);
        let graph: DiGraphMap<Mod, Edge> = builder.build().unwrap();
        assert_eq!(3, graph.node_count());
        assert_eq!(2, graph.edge_count());
        assert!(graph.contains_node(foo));
        assert!(graph.contains_node(bar));
        assert!(graph.contains_node(baz));
        assert_eq!(
            vec![(foo, bar, &Edge::Child), (bar, baz, &Edge::Child)],
            graph.all_edges().collect::<Vec<(Mod, Mod, &Edge)>>()
        );
    }

    #[test]
    #[should_panic(expected = "Module path is too long")]
    fn add_mod_panics_if_path_is_longer_than_limit() {
        let name = "foo";
        let path = std::str::from_utf8([114; MOD_PATH_SIZE - 4].as_ref()).unwrap(); // 'r'
        let concatenated = format!("{}::{}", path, name);
        assert!(concatenated.len() > MOD_PATH_SIZE);
        let mut builder = GraphBuilder::new();
        builder.add_crate_root(path);
        builder.add_mod(path, name, Visibility::Public);
    }

    #[test]
    #[should_panic(expected = "Module cannot depend on itself")]
    fn adding_a_dependency_to_the_same_module_panics() {
        let mut builder = GraphBuilder::new();
        builder.add_crate_root("root");
        builder.add_mod("root", "sub", Visibility::Public);
        builder.add_dep("root::sub", "root::sub");
    }

    #[test]
    fn adding_a_dependency_is_idempotent() {
        let mut builder = GraphBuilder::new();
        builder.add_crate_root("foo");
        builder.add_mod("foo", "bar", Visibility::Public);
        builder.add_mod("foo", "baz", Visibility::Private);
        builder.add_dep("foo::bar", "foo::baz");
        builder.add_dep("foo::bar", "foo::baz");
        let graph = builder.build().unwrap();
        assert_eq!(3, graph.edge_count());
    }

    #[test]
    fn add_dep_requires_both_modules_to_be_defined() {
        {
            let mut builder = GraphBuilder::new();
            builder.add_crate_root("foo");
            builder.add_mod("foo", "bar", Visibility::Public);
            builder.add_dep("foo::bar", "foo::baz");
            assert_eq!(
                Some(GraphError::UnknownModule(String::from("foo::baz"))),
                builder.build().err()
            );
        }
        {
            let mut builder = GraphBuilder::new();
            builder.add_crate_root("foo");
            builder.add_mod("foo", "bar", Visibility::Private);
            builder.add_dep("foo::bar", "foo::baz");
            builder.add_dep("foo::bar", "foo::fubar");
            builder.add_mod("foo", "baz", Visibility::Private);
            assert_eq!(
                Some(GraphError::UnknownModule(String::from("foo::fubar"))),
                builder.build().err()
            );
        }
    }

    #[test]
    #[should_panic(expected = "Trying to add dependency from an unknown module")]
    fn dependency_from_unknown_module_panics() {
        {
            let mut builder = GraphBuilder::new();
            builder.add_crate_root("foo");
            builder.add_dep("foo::bar", "foo::baz");
        }
    }

    #[test]
    fn dependency_can_be_added_before_dependent_module_is_added() {
        let mut builder = GraphBuilder::new();
        builder.add_crate_root("foo");
        builder.add_mod("foo", "bar", Visibility::Public);
        builder.add_dep("foo::bar", "foo::baz");
        builder.add_mod("foo", "baz", Visibility::Public);
        assert_eq!(3, builder.build().unwrap().edge_count());
    }
}
<|endoftext|>


<commit_before>use rustc::ty;
use rustc::ty::layout::Primitive;
use rustc::mir;

use super::*;

use helpers::EvalContextExt as HelperEvalContextExt;

pub trait EvalContextExt<'tcx> {
    fn ptr_op(
        &self,
        bin_op: mir::BinOp,
        left: Scalar,
        left_ty: ty::Ty<'tcx>,
        right: Scalar,
        right_ty: ty::Ty<'tcx>,
    ) -> EvalResult<'tcx, Option<(Scalar, bool)>>;

    fn ptr_int_arithmetic(
        &self,
        bin_op: mir::BinOp,
        left: Pointer,
        right: i128,
        signed: bool,
    ) -> EvalResult<'tcx, (Scalar, bool)>;
}

impl<'a, 'mir, 'tcx> EvalContextExt<'tcx> for EvalContext<'a, 'mir, 'tcx, super::Evaluator<'tcx>> {
    fn ptr_op(
        &self,
        bin_op: mir::BinOp,
        left: Scalar,
        left_ty: ty::Ty<'tcx>,
        right: Scalar,
        right_ty: ty::Ty<'tcx>,
    ) -> EvalResult<'tcx, Option<(Scalar, bool)>> {
        use rustc::mir::BinOp::*;
        use rustc::ty::layout::Integer::*;
        let usize = Primitive::Int(match self.memory.pointer_size().bytes() {
            1 => I8,
            2 => I16,
            4 => I32,
            8 => I64,
            16 => I128,
            _ => unreachable!(),
        }, false);
        let isize = Primitive::Int(match self.memory.pointer_size().bytes() {
            1 => I8,
            2 => I16,
            4 => I32,
            8 => I64,
            16 => I128,
            _ => unreachable!(),
        }, true);
        let left_kind = match self.layout_of(left_ty)?.abi {
            ty::layout::Abi::Scalar(ref scalar) => scalar.value,
            _ => Err(EvalErrorKind::TypeNotPrimitive(left_ty))?,
        };
        let right_kind = match self.layout_of(right_ty)?.abi {
            ty::layout::Abi::Scalar(ref scalar) => scalar.value,
            _ => Err(EvalErrorKind::TypeNotPrimitive(right_ty))?,
        };
        match bin_op {
            Offset if left_kind == Primitive::Pointer && right_kind == usize => {
                let pointee_ty = left_ty
                    .builtin_deref(true)
                    .expect("Offset called on non-ptr type")
                    .ty;
                let ptr = self.pointer_offset(
                    left.into(),
                    pointee_ty,
                    right.to_bits(self.memory.pointer_size())? as i64,
                )?;
                Ok(Some((ptr, false)))
            }
            // These work on anything
            Eq if left_kind == right_kind => {
                let result = match (left, right) {
                    (Scalar::Bits { .. }, Scalar::Bits { .. }) => {
                        left.to_bits(self.memory.pointer_size())? == right.to_bits(self.memory.pointer_size())?
                    },
                    (Scalar::Ptr(left), Scalar::Ptr(right)) => left == right,
                    _ => false,
                };
                Ok(Some((Scalar::from_bool(result), false)))
            }
            Ne if left_kind == right_kind => {
                let result = match (left, right) {
                    (Scalar::Bits { .. }, Scalar::Bits { .. }) => {
                        left.to_bits(self.memory.pointer_size())? != right.to_bits(self.memory.pointer_size())?
                    },
                    (Scalar::Ptr(left), Scalar::Ptr(right)) => left != right,
                    _ => true,
                };
                Ok(Some((Scalar::from_bool(result), false)))
            }
            // These need both pointers to be in the same allocation
            Lt | Le | Gt | Ge | Sub
                if left_kind == right_kind &&
                       (left_kind == Primitive::Pointer || left_kind == usize || left_kind == isize) &&
                       left.is_ptr() && right.is_ptr() => {
                let left = left.to_ptr()?;
                let right = right.to_ptr()?;
                if left.alloc_id == right.alloc_id {
                    let res = match bin_op {
                        Lt => left.offset < right.offset,
                        Le => left.offset <= right.offset,
                        Gt => left.offset > right.offset,
                        Ge => left.offset >= right.offset,
                        Sub => {
                            return self.binary_op(
                                Sub,
                                Scalar::Bits { bits: left.offset.bytes() as u128, defined: self.memory.pointer_size().bits() as u8 },
                                self.tcx.types.usize,
                                Scalar::Bits { bits: right.offset.bytes() as u128, defined: self.memory.pointer_size().bits() as u8 },
                                self.tcx.types.usize,
                            ).map(Some)
                        }
                        _ => bug!("We already established it has to be one of these operators."),
                    };
                    Ok(Some((Scalar::from_bool(res), false)))
                } else {
                    // Both are pointers, but from different allocations.
                    err!(InvalidPointerMath)
                }
            }
            // These work if one operand is a pointer, the other an integer
            Add | BitAnd | Sub
                if left_kind == right_kind && (left_kind == usize || left_kind == isize) &&
                       left.is_ptr() && right.is_bits() => {
                // Cast to i128 is fine as we checked the kind to be ptr-sized
                self.ptr_int_arithmetic(
                    bin_op,
                    left.to_ptr()?,
                    right.to_bits(self.memory.pointer_size())? as i128,
                    left_kind == isize,
                ).map(Some)
            }
            Add | BitAnd
                if left_kind == right_kind && (left_kind == usize || left_kind == isize) &&
                       left.is_bits() && right.is_ptr() => {
                // This is a commutative operation, just swap the operands
                self.ptr_int_arithmetic(
                    bin_op,
                    right.to_ptr()?,
                    left.to_bits(self.memory.pointer_size())? as i128,
                    left_kind == isize,
                ).map(Some)
            }
            _ => Ok(None),
        }
    }

    fn ptr_int_arithmetic(
        &self,
        bin_op: mir::BinOp,
        left: Pointer,
        right: i128,
        signed: bool,
    ) -> EvalResult<'tcx, (Scalar, bool)> {
        use rustc::mir::BinOp::*;

        fn map_to_primval((res, over): (Pointer, bool)) -> (Scalar, bool) {
            (Scalar::Ptr(res), over)
        }

        Ok(match bin_op {
            Sub =>
                // The only way this can overflow is by underflowing, so signdeness of the right operands does not matter
                map_to_primval(left.overflowing_signed_offset(-right, self)),
            Add if signed =>
                map_to_primval(left.overflowing_signed_offset(right, self)),
            Add if !signed =>
                map_to_primval(left.overflowing_offset(Size::from_bytes(right as u64), self)),

            BitAnd if !signed => {
                let base_mask : u64 = !(self.memory.get(left.alloc_id)?.align.abi() - 1);
                let right = right as u64;
                if right & base_mask == base_mask {
                    // Case 1: The base address bits are all preserved, i.e., right is all-1 there
                    (Scalar::Ptr(Pointer::new(left.alloc_id, Size::from_bytes(left.offset.bytes() & right))), false)
                } else if right & base_mask == 0 {
                    // Case 2: The base address bits are all taken away, i.e., right is all-0 there
                    (Scalar::Bits { bits: (left.offset.bytes() & right) as u128, defined: 128 }, false)
                } else {
                    return err!(ReadPointerAsBytes);
                }
            }

            _ => {
                let msg = format!("unimplemented binary op on pointer {:?}: {:?}, {:?} ({})", bin_op, left, right, if signed { "signed" } else { "unsigned" });
                return err!(Unimplemented(msg));
            }
        })
    }
}
<commit_msg>Comparing non-pointer-size types should be possible<commit_after>use rustc::ty;
use rustc::ty::layout::Primitive;
use rustc::mir;

use super::*;

use helpers::EvalContextExt as HelperEvalContextExt;

pub trait EvalContextExt<'tcx> {
    fn ptr_op(
        &self,
        bin_op: mir::BinOp,
        left: Scalar,
        left_ty: ty::Ty<'tcx>,
        right: Scalar,
        right_ty: ty::Ty<'tcx>,
    ) -> EvalResult<'tcx, Option<(Scalar, bool)>>;

    fn ptr_int_arithmetic(
        &self,
        bin_op: mir::BinOp,
        left: Pointer,
        right: i128,
        signed: bool,
    ) -> EvalResult<'tcx, (Scalar, bool)>;
}

impl<'a, 'mir, 'tcx> EvalContextExt<'tcx> for EvalContext<'a, 'mir, 'tcx, super::Evaluator<'tcx>> {
    fn ptr_op(
        &self,
        bin_op: mir::BinOp,
        left: Scalar,
        left_ty: ty::Ty<'tcx>,
        right: Scalar,
        right_ty: ty::Ty<'tcx>,
    ) -> EvalResult<'tcx, Option<(Scalar, bool)>> {
        use rustc::mir::BinOp::*;
        use rustc::ty::layout::Integer::*;
        let usize = Primitive::Int(match self.memory.pointer_size().bytes() {
            1 => I8,
            2 => I16,
            4 => I32,
            8 => I64,
            16 => I128,
            _ => unreachable!(),
        }, false);
        let isize = Primitive::Int(match self.memory.pointer_size().bytes() {
            1 => I8,
            2 => I16,
            4 => I32,
            8 => I64,
            16 => I128,
            _ => unreachable!(),
        }, true);
        let left_layout = self.layout_of(left_ty)?;
        let left_kind = match left_layout.abi {
            ty::layout::Abi::Scalar(ref scalar) => scalar.value,
            _ => Err(EvalErrorKind::TypeNotPrimitive(left_ty))?,
        };
        let right_layout = self.layout_of(right_ty)?;
        let right_kind = match right_layout.abi {
            ty::layout::Abi::Scalar(ref scalar) => scalar.value,
            _ => Err(EvalErrorKind::TypeNotPrimitive(right_ty))?,
        };
        match bin_op {
            Offset if left_kind == Primitive::Pointer && right_kind == usize => {
                let pointee_ty = left_ty
                    .builtin_deref(true)
                    .expect("Offset called on non-ptr type")
                    .ty;
                let ptr = self.pointer_offset(
                    left.into(),
                    pointee_ty,
                    right.to_bits(self.memory.pointer_size())? as i64,
                )?;
                Ok(Some((ptr, false)))
            }
            // These work on anything
            Eq if left_kind == right_kind => {
                let result = match (left, right) {
                    (Scalar::Bits { .. }, Scalar::Bits { .. }) => {
                        left.to_bits(left_layout.size)? == right.to_bits(right_layout.size)?
                    },
                    (Scalar::Ptr(left), Scalar::Ptr(right)) => left == right,
                    _ => false,
                };
                Ok(Some((Scalar::from_bool(result), false)))
            }
            Ne if left_kind == right_kind => {
                let result = match (left, right) {
                    (Scalar::Bits { .. }, Scalar::Bits { .. }) => {
                        left.to_bits(left_layout.size)? != right.to_bits(right_layout.size)?
                    },
                    (Scalar::Ptr(left), Scalar::Ptr(right)) => left != right,
                    _ => true,
                };
                Ok(Some((Scalar::from_bool(result), false)))
            }
            // These need both pointers to be in the same allocation
            Lt | Le | Gt | Ge | Sub
                if left_kind == right_kind &&
                       (left_kind == Primitive::Pointer || left_kind == usize || left_kind == isize) &&
                       left.is_ptr() && right.is_ptr() => {
                let left = left.to_ptr()?;
                let right = right.to_ptr()?;
                if left.alloc_id == right.alloc_id {
                    let res = match bin_op {
                        Lt => left.offset < right.offset,
                        Le => left.offset <= right.offset,
                        Gt => left.offset > right.offset,
                        Ge => left.offset >= right.offset,
                        Sub => {
                            return self.binary_op(
                                Sub,
                                Scalar::Bits { bits: left.offset.bytes() as u128, defined: self.memory.pointer_size().bits() as u8 },
                                self.tcx.types.usize,
                                Scalar::Bits { bits: right.offset.bytes() as u128, defined: self.memory.pointer_size().bits() as u8 },
                                self.tcx.types.usize,
                            ).map(Some)
                        }
                        _ => bug!("We already established it has to be one of these operators."),
                    };
                    Ok(Some((Scalar::from_bool(res), false)))
                } else {
                    // Both are pointers, but from different allocations.
                    err!(InvalidPointerMath)
                }
            }
            // These work if one operand is a pointer, the other an integer
            Add | BitAnd | Sub
                if left_kind == right_kind && (left_kind == usize || left_kind == isize) &&
                       left.is_ptr() && right.is_bits() => {
                // Cast to i128 is fine as we checked the kind to be ptr-sized
                self.ptr_int_arithmetic(
                    bin_op,
                    left.to_ptr()?,
                    right.to_bits(self.memory.pointer_size())? as i128,
                    left_kind == isize,
                ).map(Some)
            }
            Add | BitAnd
                if left_kind == right_kind && (left_kind == usize || left_kind == isize) &&
                       left.is_bits() && right.is_ptr() => {
                // This is a commutative operation, just swap the operands
                self.ptr_int_arithmetic(
                    bin_op,
                    right.to_ptr()?,
                    left.to_bits(self.memory.pointer_size())? as i128,
                    left_kind == isize,
                ).map(Some)
            }
            _ => Ok(None),
        }
    }

    fn ptr_int_arithmetic(
        &self,
        bin_op: mir::BinOp,
        left: Pointer,
        right: i128,
        signed: bool,
    ) -> EvalResult<'tcx, (Scalar, bool)> {
        use rustc::mir::BinOp::*;

        fn map_to_primval((res, over): (Pointer, bool)) -> (Scalar, bool) {
            (Scalar::Ptr(res), over)
        }

        Ok(match bin_op {
            Sub =>
                // The only way this can overflow is by underflowing, so signdeness of the right operands does not matter
                map_to_primval(left.overflowing_signed_offset(-right, self)),
            Add if signed =>
                map_to_primval(left.overflowing_signed_offset(right, self)),
            Add if !signed =>
                map_to_primval(left.overflowing_offset(Size::from_bytes(right as u64), self)),

            BitAnd if !signed => {
                let base_mask : u64 = !(self.memory.get(left.alloc_id)?.align.abi() - 1);
                let right = right as u64;
                if right & base_mask == base_mask {
                    // Case 1: The base address bits are all preserved, i.e., right is all-1 there
                    (Scalar::Ptr(Pointer::new(left.alloc_id, Size::from_bytes(left.offset.bytes() & right))), false)
                } else if right & base_mask == 0 {
                    // Case 2: The base address bits are all taken away, i.e., right is all-0 there
                    (Scalar::Bits { bits: (left.offset.bytes() & right) as u128, defined: 128 }, false)
                } else {
                    return err!(ReadPointerAsBytes);
                }
            }

            _ => {
                let msg = format!("unimplemented binary op on pointer {:?}: {:?}, {:?} ({})", bin_op, left, right, if signed { "signed" } else { "unsigned" });
                return err!(Unimplemented(msg));
            }
        })
    }
}
<|endoftext|>