<commit_before>use std::borrow::Cow;
use std::time::{Duration, Instant};

use regex::{Captures, Regex};

use console::{measure_text_width, Style};

pub fn duration_to_secs(d: Duration) -> f64 {
    d.as_secs() as f64 + d.subsec_nanos() as f64 / 1_000_000_000f64
}

pub fn secs_to_duration(s: f64) -> Duration {
    let secs = s.trunc() as u64;
    let nanos = (s.fract() * 1_000_000_000f64) as u32;
    Duration::new(secs, nanos)
}

pub struct Estimate {
    buf: Vec<f64>,
    buf_cap: usize,
    last_idx: usize,
    started: Option<(Instant, u64)>,
}

impl Estimate {
    pub fn new() -> Estimate {
        Estimate {
            buf: vec![],
            buf_cap: 10,
            last_idx: 0,
            started: None,
        }
    }

    pub fn reset(&mut self) {
        self.buf.clear();
        self.last_idx = 0;
        self.started = None;
    }

    pub fn record_step(&mut self, value: u64) {
        // record initial position
        let (started_time, started_value) = match self.started {
            None => {
                let rv = (Instant::now(), value);
                self.started = Some(rv);
                rv
            }
            Some(value) => value,
        };

        let item = if value == 0 {
            0.0
        } else {
            let divisor = value.saturating_sub(started_value) as f64;
            if divisor == 0.0 {
                0.0
            } else {
                duration_to_secs(started_time.elapsed()) / divisor
            }
        };
        if self.buf.len() >= self.buf_cap {
            let idx = self.last_idx % self.buf.len();
            self.buf[idx] = item;
        } else {
            self.buf.push(item);
        }
        self.last_idx += 1;
    }

    pub fn time_per_step(&self) -> Duration {
        if self.buf.is_empty() {
            Duration::new(0, 0)
        } else {
            secs_to_duration(self.buf.iter().sum::<f64>() / self.buf.len() as f64)
        }
    }
}

#[derive(PartialEq, Eq, Debug, Copy, Clone)]
pub enum Alignment {
    Left,
    Center,
    Right,
}

#[derive(Debug)]
pub struct TemplateVar<'a> {
    pub key: &'a str,
    pub align: Alignment,
    pub truncate: bool,
    pub width: Option<usize>,
    pub style: Option<Style>,
    pub alt_style: Option<Style>,
    pub last_element: bool,
}

impl<'a> TemplateVar<'a> {
    pub fn duplicate_for_key<'b>(&self, key: &'b str) -> TemplateVar<'b> {
        TemplateVar {
            key: key,
            align: self.align,
            truncate: self.truncate,
            width: self.width,
            style: self.style.clone(),
            alt_style: self.alt_style.clone(),
            last_element: self.last_element,
        }
    }
}

pub fn expand_template<'a, F: Fn(&TemplateVar) -> String>(s: &'a str, f: F) -> Cow<'a, str> {
    lazy_static! {
        static ref VAR_RE: Regex = Regex::new(r"(\}\})|\{(\{|[^}]+\})").unwrap();
        static ref KEY_RE: Regex = Regex::new(
            r"(?x)
                ([^:]+)
                (?:
                    :
                    ([<^>])?
                    (\d+)?
                    (!)?
                    (?:\.([a-z_]+(?:\.[a-z_]+)*))?
                    (?:/([a-z_]+(?:\.[a-z_]+)*))?
                )?
            "
        ).unwrap();
    }
    VAR_RE.replace_all(s, |caps: &Captures| {
        if caps.get(1).is_some() {
            return "}".into();
        }
        let key = &caps[2];
        if key == "{" {
            return "{".into();
        }
        let mut var = TemplateVar {
            key: key,
            align: Alignment::Left,
            truncate: false,
            width: None,
            style: None,
            alt_style: None,
            last_element: caps.get(0).unwrap().end() >= s.len(),
        };
        if let Some(opt_caps) = KEY_RE.captures(&key[..key.len() - 1]) {
            if let Some(short_key) = opt_caps.get(1) {
                var.key = short_key.as_str();
            }
            var.align = match opt_caps.get(2).map(|x| x.as_str()) {
                Some("<") => Alignment::Left,
                Some("^") => Alignment::Center,
                Some(">") => Alignment::Right,
                _ => Alignment::Left,
            };
            if let Some(width) = opt_caps.get(3) {
                var.width = Some(width.as_str().parse().unwrap());
            }
            if opt_caps.get(4).is_some() {
                var.truncate = true;
            }
            if let Some(style) = opt_caps.get(5) {
                var.style = Some(Style::from_dotted_str(style.as_str()));
            }
            if let Some(alt_style) = opt_caps.get(6) {
                var.alt_style = Some(Style::from_dotted_str(alt_style.as_str()));
            }
        }
        let mut rv = f(&var);
        if let Some(width) = var.width {
            rv = pad_str(&rv, width, var.align, var.truncate).to_string()
        }
        if let Some(s) = var.style {
            rv = s.apply_to(rv).to_string();
        }
        rv
    })
}

pub fn pad_str<'a>(s: &'a str, width: usize, align: Alignment, truncate: bool) -> Cow<'a, str> {
    let cols = measure_text_width(s);

    if cols >= width {
        return if truncate {
            Cow::Borrowed(s.get(..width).unwrap_or(s))
        } else {
            Cow::Borrowed(s)
        };
    }

    let diff = width.saturating_sub(cols);

    let (left_pad, right_pad) = match align {
        Alignment::Left => (0, diff),
        Alignment::Right => (diff, 0),
        Alignment::Center => (diff / 2, diff.saturating_sub(diff / 2)),
    };

    let mut rv = String::new();
    for _ in 0..left_pad {
        rv.push(' ');
    }
    rv.push_str(s);
    for _ in 0..right_pad {
        rv.push(' ');
    }
    Cow::Owned(rv)
}

#[test]
fn test_expand_template() {
    let rv = expand_template("{{ {foo} {bar} }}", |var| var.key.to_uppercase());
    assert_eq!(&rv, "{ FOO BAR }");
}

#[test]
fn test_expand_template_flags() {
    use console::set_colors_enabled;
    set_colors_enabled(true);

    let rv = expand_template("{foo:5}", |var| {
        assert_eq!(var.key, "foo");
        assert_eq!(var.width, Some(5));
        "XXX".into()
    });
    assert_eq!(&rv, "XXX  ");

    let rv = expand_template("{foo:.red.on_blue}", |var| {
        assert_eq!(var.key, "foo");
        assert_eq!(var.width, None);
        assert_eq!(var.align, Alignment::Left);
        assert_eq!(var.style, Some(Style::new().red().on_blue()));
        "XXX".into()
    });
    assert_eq!(&rv, "\u{1b}[31m\u{1b}[44mXXX\u{1b}[0m");

    let rv = expand_template("{foo:^5.red.on_blue}", |var| {
        assert_eq!(var.key, "foo");
        assert_eq!(var.width, Some(5));
        assert_eq!(var.align, Alignment::Center);
        assert_eq!(var.style, Some(Style::new().red().on_blue()));
        "XXX".into()
    });
    assert_eq!(&rv, "\u{1b}[31m\u{1b}[44m XXX \u{1b}[0m");

    let rv = expand_template("{foo:^5.red.on_blue/green.on_cyan}", |var| {
        assert_eq!(var.key, "foo");
        assert_eq!(var.width, Some(5));
        assert_eq!(var.align, Alignment::Center);
        assert_eq!(var.style, Some(Style::new().red().on_blue()));
        assert_eq!(var.alt_style, Some(Style::new().green().on_cyan()));
        "XXX".into()
    });
    assert_eq!(&rv, "\u{1b}[31m\u{1b}[44m XXX \u{1b}[0m");
}

#[test]
fn test_duration_stuff() {
    let duration = Duration::new(42, 100_000_000);
    let secs = duration_to_secs(duration);
    assert_eq!(secs_to_duration(secs), duration);
}
<commit_msg>cargo fmt<commit_after>use std::borrow::Cow;
use std::time::{Duration, Instant};

use regex::{Captures, Regex};

use console::{measure_text_width, Style};

pub fn duration_to_secs(d: Duration) -> f64 {
    d.as_secs() as f64 + d.subsec_nanos() as f64 / 1_000_000_000f64
}

pub fn secs_to_duration(s: f64) -> Duration {
    let secs = s.trunc() as u64;
    let nanos = (s.fract() * 1_000_000_000f64) as u32;
    Duration::new(secs, nanos)
}

pub struct Estimate {
    buf: Vec<f64>,
    buf_cap: usize,
    last_idx: usize,
    started: Option<(Instant, u64)>,
}

impl Estimate {
    pub fn new() -> Estimate {
        Estimate {
            buf: vec![],
            buf_cap: 10,
            last_idx: 0,
            started: None,
        }
    }

    pub fn reset(&mut self) {
        self.buf.clear();
        self.last_idx = 0;
        self.started = None;
    }

    pub fn record_step(&mut self, value: u64) {
        // record initial position
        let (started_time, started_value) = match self.started {
            None => {
                let rv = (Instant::now(), value);
                self.started = Some(rv);
                rv
            }
            Some(value) => value,
        };

        let item = if value == 0 {
            0.0
        } else {
            let divisor = value.saturating_sub(started_value) as f64;
            if divisor == 0.0 {
                0.0
            } else {
                duration_to_secs(started_time.elapsed()) / divisor
            }
        };
        if self.buf.len() >= self.buf_cap {
            let idx = self.last_idx % self.buf.len();
            self.buf[idx] = item;
        } else {
            self.buf.push(item);
        }
        self.last_idx += 1;
    }

    pub fn time_per_step(&self) -> Duration {
        if self.buf.is_empty() {
            Duration::new(0, 0)
        } else {
            secs_to_duration(self.buf.iter().sum::<f64>() / self.buf.len() as f64)
        }
    }
}

#[derive(PartialEq, Eq, Debug, Copy, Clone)]
pub enum Alignment {
    Left,
    Center,
    Right,
}

#[derive(Debug)]
pub struct TemplateVar<'a> {
    pub key: &'a str,
    pub align: Alignment,
    pub truncate: bool,
    pub width: Option<usize>,
    pub style: Option<Style>,
    pub alt_style: Option<Style>,
    pub last_element: bool,
}

impl<'a> TemplateVar<'a> {
    pub fn duplicate_for_key<'b>(&self, key: &'b str) -> TemplateVar<'b> {
        TemplateVar {
            key: key,
            align: self.align,
            truncate: self.truncate,
            width: self.width,
            style: self.style.clone(),
            alt_style: self.alt_style.clone(),
            last_element: self.last_element,
        }
    }
}

pub fn expand_template<'a, F: Fn(&TemplateVar) -> String>(s: &'a str, f: F) -> Cow<'a, str> {
    lazy_static! {
        static ref VAR_RE: Regex = Regex::new(r"(\}\})|\{(\{|[^}]+\})").unwrap();
        static ref KEY_RE: Regex = Regex::new(
            r"(?x)
                ([^:]+)
                (?:
                    :
                    ([<^>])?
                    (\d+)?
                    (!)?
                    (?:\.([a-z_]+(?:\.[a-z_]+)*))?
                    (?:/([a-z_]+(?:\.[a-z_]+)*))?
                )?
            "
        )
        .unwrap();
    }
    VAR_RE.replace_all(s, |caps: &Captures| {
        if caps.get(1).is_some() {
            return "}".into();
        }
        let key = &caps[2];
        if key == "{" {
            return "{".into();
        }
        let mut var = TemplateVar {
            key: key,
            align: Alignment::Left,
            truncate: false,
            width: None,
            style: None,
            alt_style: None,
            last_element: caps.get(0).unwrap().end() >= s.len(),
        };
        if let Some(opt_caps) = KEY_RE.captures(&key[..key.len() - 1]) {
            if let Some(short_key) = opt_caps.get(1) {
                var.key = short_key.as_str();
            }
            var.align = match opt_caps.get(2).map(|x| x.as_str()) {
                Some("<") => Alignment::Left,
                Some("^") => Alignment::Center,
                Some(">") => Alignment::Right,
                _ => Alignment::Left,
            };
            if let Some(width) = opt_caps.get(3) {
                var.width = Some(width.as_str().parse().unwrap());
            }
            if opt_caps.get(4).is_some() {
                var.truncate = true;
            }
            if let Some(style) = opt_caps.get(5) {
                var.style = Some(Style::from_dotted_str(style.as_str()));
            }
            if let Some(alt_style) = opt_caps.get(6) {
                var.alt_style = Some(Style::from_dotted_str(alt_style.as_str()));
            }
        }
        let mut rv = f(&var);
        if let Some(width) = var.width {
            rv = pad_str(&rv, width, var.align, var.truncate).to_string()
        }
        if let Some(s) = var.style {
            rv = s.apply_to(rv).to_string();
        }
        rv
    })
}

pub fn pad_str<'a>(s: &'a str, width: usize, align: Alignment, truncate: bool) -> Cow<'a, str> {
    let cols = measure_text_width(s);

    if cols >= width {
        return if truncate {
            Cow::Borrowed(s.get(..width).unwrap_or(s))
        } else {
            Cow::Borrowed(s)
        };
    }

    let diff = width.saturating_sub(cols);

    let (left_pad, right_pad) = match align {
        Alignment::Left => (0, diff),
        Alignment::Right => (diff, 0),
        Alignment::Center => (diff / 2, diff.saturating_sub(diff / 2)),
    };

    let mut rv = String::new();
    for _ in 0..left_pad {
        rv.push(' ');
    }
    rv.push_str(s);
    for _ in 0..right_pad {
        rv.push(' ');
    }
    Cow::Owned(rv)
}

#[test]
fn test_expand_template() {
    let rv = expand_template("{{ {foo} {bar} }}", |var| var.key.to_uppercase());
    assert_eq!(&rv, "{ FOO BAR }");
}

#[test]
fn test_expand_template_flags() {
    use console::set_colors_enabled;
    set_colors_enabled(true);

    let rv = expand_template("{foo:5}", |var| {
        assert_eq!(var.key, "foo");
        assert_eq!(var.width, Some(5));
        "XXX".into()
    });
    assert_eq!(&rv, "XXX  ");

    let rv = expand_template("{foo:.red.on_blue}", |var| {
        assert_eq!(var.key, "foo");
        assert_eq!(var.width, None);
        assert_eq!(var.align, Alignment::Left);
        assert_eq!(var.style, Some(Style::new().red().on_blue()));
        "XXX".into()
    });
    assert_eq!(&rv, "\u{1b}[31m\u{1b}[44mXXX\u{1b}[0m");

    let rv = expand_template("{foo:^5.red.on_blue}", |var| {
        assert_eq!(var.key, "foo");
        assert_eq!(var.width, Some(5));
        assert_eq!(var.align, Alignment::Center);
        assert_eq!(var.style, Some(Style::new().red().on_blue()));
        "XXX".into()
    });
    assert_eq!(&rv, "\u{1b}[31m\u{1b}[44m XXX \u{1b}[0m");

    let rv = expand_template("{foo:^5.red.on_blue/green.on_cyan}", |var| {
        assert_eq!(var.key, "foo");
        assert_eq!(var.width, Some(5));
        assert_eq!(var.align, Alignment::Center);
        assert_eq!(var.style, Some(Style::new().red().on_blue()));
        assert_eq!(var.alt_style, Some(Style::new().green().on_cyan()));
        "XXX".into()
    });
    assert_eq!(&rv, "\u{1b}[31m\u{1b}[44m XXX \u{1b}[0m");
}

#[test]
fn test_duration_stuff() {
    let duration = Duration::new(42, 100_000_000);
    let secs = duration_to_secs(duration);
    assert_eq!(secs_to_duration(secs), duration);
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    ChainInfo, FullNode, HealthCheckError, LocalNode, LocalVersion, Node, Swarm, SwarmExt,
    Validator, Version,
};
use anyhow::{anyhow, Result};
use diem_config::config::NodeConfig;
use diem_genesis_tool::{fullnode_builder::FullnodeConfig, validator_builder::ValidatorBuilder};
use diem_sdk::{
    crypto::ed25519::Ed25519PrivateKey,
    types::{
        chain_id::ChainId, transaction::Transaction, waypoint::Waypoint, AccountKey, LocalAccount,
        PeerId,
    },
};
use std::{
    collections::HashMap,
    convert::TryFrom,
    fs, mem,
    num::NonZeroUsize,
    ops,
    path::{Path, PathBuf},
    sync::Arc,
    time::{Duration, Instant},
};
use tempfile::TempDir;

#[derive(Debug)]
pub enum SwarmDirectory {
    Persistent(PathBuf),
    Temporary(TempDir),
}

impl SwarmDirectory {
    pub fn persist(&mut self) {
        match self {
            SwarmDirectory::Persistent(_) => {}
            SwarmDirectory::Temporary(_) => {
                let mut temp = SwarmDirectory::Persistent(PathBuf::new());
                mem::swap(self, &mut temp);
                let _ = mem::replace(self, temp.into_persistent());
            }
        }
    }

    pub fn into_persistent(self) -> Self {
        match self {
            SwarmDirectory::Temporary(tempdir) => SwarmDirectory::Persistent(tempdir.into_path()),
            SwarmDirectory::Persistent(dir) => SwarmDirectory::Persistent(dir),
        }
    }
}

impl ops::Deref for SwarmDirectory {
    type Target = Path;

    fn deref(&self) -> &Self::Target {
        match self {
            SwarmDirectory::Persistent(dir) => dir.deref(),
            SwarmDirectory::Temporary(dir) => dir.path(),
        }
    }
}

impl AsRef<Path> for SwarmDirectory {
    fn as_ref(&self) -> &Path {
        match self {
            SwarmDirectory::Persistent(dir) => dir.as_ref(),
            SwarmDirectory::Temporary(dir) => dir.as_ref(),
        }
    }
}

#[derive(Debug)]
struct ValidatorNetworkAddressEncryptionKey {
    key: diem_sdk::types::network_address::encrypted::Key,
    version: diem_sdk::types::network_address::encrypted::KeyVersion,
}

pub struct LocalSwarmBuilder {
    versions: Arc<HashMap<Version, LocalVersion>>,
    initial_version: Option<Version>,
    template: NodeConfig,
    number_of_validators: NonZeroUsize,
    dir: Option<PathBuf>,
}

impl LocalSwarmBuilder {
    pub fn new(versions: Arc<HashMap<Version, LocalVersion>>) -> Self {
        Self {
            versions,
            initial_version: None,
            template: NodeConfig::default_for_validator(),
            number_of_validators: NonZeroUsize::new(1).unwrap(),
            dir: None,
        }
    }

    pub fn initial_version(mut self, initial_version: Version) -> Self {
        self.initial_version = Some(initial_version);
        self
    }

    pub fn template(mut self, template: NodeConfig) -> Self {
        self.template = template;
        self
    }

    pub fn number_of_validators(mut self, number_of_validators: NonZeroUsize) -> Self {
        self.number_of_validators = number_of_validators;
        self
    }

    pub fn dir<T: AsRef<Path>>(mut self, dir: T) -> Self {
        self.dir = Some(dir.as_ref().into());
        self
    }

    pub fn build<R>(self, rng: R) -> Result<LocalSwarm>
    where
        R: ::rand::RngCore + ::rand::CryptoRng,
    {
        let dir = if let Some(dir) = self.dir {
            if dir.exists() {
                fs::remove_dir_all(&dir)?;
            }
            fs::create_dir_all(&dir)?;
            SwarmDirectory::Persistent(dir)
        } else {
            SwarmDirectory::Temporary(TempDir::new()?)
        };

        let (root_keys, genesis, genesis_waypoint, validators) = ValidatorBuilder::new(
            &dir,
            diem_framework_releases::current_module_blobs().to_vec(),
        )
        .num_validators(self.number_of_validators)
        .template(self.template)
        .build(rng)?;

        // Get the initial version to start the nodes with, either the one provided or fallback to
        // using the the latest version
        let versions = self.versions;
        let initial_version = self.initial_version.unwrap_or_else(|| {
            versions
                .iter()
                .max_by(|v1, v2| v1.0.cmp(v2.0))
                .unwrap()
                .0
                .clone()
        });
        let version = versions.get(&initial_version).unwrap();

        let validators = validators
            .into_iter()
            .map(|v| {
                let node = LocalNode::new(version.to_owned(), v.name, v.directory)?;
                Ok((node.peer_id(), node))
            })
            .collect::<Result<HashMap<_, _>>>()?;

        let validator_network_address_encryption_key = ValidatorNetworkAddressEncryptionKey {
            key: root_keys.validator_network_address_encryption_key,
            version: root_keys.validator_network_address_encryption_key_version,
        };

        let root_account = LocalAccount::new(
            diem_sdk::types::account_config::diem_root_address(),
            AccountKey::from_private_key(root_keys.root_key),
            1,
        );

        // Designated dealer account is built to have the same private key as the TC account
        let designated_dealer_account = LocalAccount::new(
            diem_sdk::types::account_config::testnet_dd_account_address(),
            AccountKey::from_private_key(
                Ed25519PrivateKey::try_from(root_keys.treasury_compliance_key.to_bytes().as_ref())
                    .unwrap(),
            ),
            0,
        );

        let treasury_compliance_account = LocalAccount::new(
            diem_sdk::types::account_config::treasury_compliance_account_address(),
            AccountKey::from_private_key(root_keys.treasury_compliance_key),
            0,
        );

        Ok(LocalSwarm {
            node_name_counter: validators.len() as u64,
            genesis,
            genesis_waypoint,
            versions,
            validators,
            fullnodes: HashMap::new(),
            dir,
            validator_network_address_encryption_key,
            root_account,
            treasury_compliance_account,
            designated_dealer_account,
            chain_id: ChainId::test(),
        })
    }
}

#[derive(Debug)]
pub struct LocalSwarm {
    node_name_counter: u64,
    genesis: Transaction,
    genesis_waypoint: Waypoint,
    versions: Arc<HashMap<Version, LocalVersion>>,
    validators: HashMap<PeerId, LocalNode>,
    fullnodes: HashMap<PeerId, LocalNode>,
    dir: SwarmDirectory,
    validator_network_address_encryption_key: ValidatorNetworkAddressEncryptionKey,
    root_account: LocalAccount,
    treasury_compliance_account: LocalAccount,
    designated_dealer_account: LocalAccount,
    chain_id: ChainId,
}

impl LocalSwarm {
    pub fn builder(versions: Arc<HashMap<Version, LocalVersion>>) -> LocalSwarmBuilder {
        LocalSwarmBuilder::new(versions)
    }

    pub fn launch(&mut self) -> Result<()> {
        // Start all the validators
        for validator in self.validators.values_mut() {
            validator.start()?;
        }

        // Wait for all of them to startup
        let deadline = Instant::now() + Duration::from_secs(60);
        self.wait_for_startup()?;
        self.wait_for_connectivity(deadline)?;
        self.liveness_check(deadline)?;

        Ok(())
    }

    fn wait_for_startup(&mut self) -> Result<()> {
        let num_attempts = 10;
        let mut done = vec![false; self.validators.len()];
        for i in 0..num_attempts {
            println!("Wait for startup attempt: {} of {}", i, num_attempts);
            for (node, done) in self.validators.values_mut().zip(done.iter_mut()) {
                if *done {
                    continue;
                }
                match node.health_check() {
                    Ok(()) => *done = true,

                    Err(HealthCheckError::Unknown(e)) => {
                        return Err(anyhow!("Diem node '{}' is not running: {}", node.name(), e));
                    }
                    Err(HealthCheckError::NotRunning) => {
                        return Err(anyhow!("Diem node '{}' is not running", node.name()));
                    }
                    Err(HealthCheckError::RpcFailure(e)) => {
                        println!("rpc error: {}", e);
                        continue;
                    }
                }
            }

            // Check if all the nodes have been successfully launched
            if done.iter().all(|status| *status) {
                return Ok(());
            }

            ::std::thread::sleep(::std::time::Duration::from_millis(1000));
        }

        Err(anyhow!("Launching Swarm timed out"))
    }

    fn add_fullnode(&mut self, version: &Version, template: NodeConfig) -> Result<PeerId> {
        let name = self.node_name_counter.to_string();
        self.node_name_counter += 1;
        let fullnode_config = FullnodeConfig::public_fullnode(
            name,
            self.dir.as_ref(),
            template,
            &self.genesis_waypoint,
            &self.genesis,
        )?;

        let version = self.versions.get(version).unwrap();
        let mut fullnode = LocalNode::new(
            version.to_owned(),
            fullnode_config.name,
            fullnode_config.directory,
        )?;

        let peer_id = fullnode.peer_id();
        fullnode.start()?;

        self.fullnodes.insert(peer_id, fullnode);

        Ok(peer_id)
    }
}

impl Swarm for LocalSwarm {
    fn health_check(&mut self) -> Result<()> {
        todo!()
    }

    fn validators<'a>(&'a self) -> Box<dyn Iterator<Item = &'a dyn Validator> + 'a> {
        Box::new(self.validators.values().map(|v| v as &'a dyn Validator))
    }

    fn validators_mut<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a mut dyn Validator> + 'a> {
        Box::new(
            self.validators
                .values_mut()
                .map(|v| v as &'a mut dyn Validator),
        )
    }

    fn validator(&self, id: PeerId) -> Option<&dyn Validator> {
        self.validators.get(&id).map(|v| v as &dyn Validator)
    }

    fn validator_mut(&mut self, id: PeerId) -> Option<&mut dyn Validator> {
        self.validators
            .get_mut(&id)
            .map(|v| v as &mut dyn Validator)
    }

    fn upgrade_validator(&mut self, id: PeerId, version: &Version) -> Result<()> {
        let version = self
            .versions
            .get(version)
            .cloned()
            .ok_or_else(|| anyhow!("Invalid version: {:?}", version))?;
        let validator = self
            .validators
            .get_mut(&id)
            .ok_or_else(|| anyhow!("Invalid id: {}", id))?;
        validator.upgrade(version)
    }

    fn full_nodes<'a>(&'a self) -> Box<dyn Iterator<Item = &'a dyn FullNode> + 'a> {
        Box::new(self.fullnodes.values().map(|v| v as &'a dyn FullNode))
    }

    fn full_nodes_mut<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a mut dyn FullNode> + 'a> {
        Box::new(
            self.fullnodes
                .values_mut()
                .map(|v| v as &'a mut dyn FullNode),
        )
    }

    fn full_node(&self, id: PeerId) -> Option<&dyn FullNode> {
        self.fullnodes.get(&id).map(|v| v as &dyn FullNode)
    }

    fn full_node_mut(&mut self, id: PeerId) -> Option<&mut dyn FullNode> {
        self.fullnodes.get_mut(&id).map(|v| v as &mut dyn FullNode)
    }

    fn add_validator(&mut self, _version: &Version, _template: NodeConfig) -> Result<PeerId> {
        todo!()
    }

    fn remove_validator(&mut self, _id: PeerId) -> Result<()> {
        todo!()
    }

    fn add_full_node(&mut self, version: &Version, template: NodeConfig) -> Result<PeerId> {
        self.add_fullnode(version, template)
    }

    fn remove_full_node(&mut self, id: PeerId) -> Result<()> {
        if let Some(mut fullnode) = self.fullnodes.remove(&id) {
            fullnode.stop();
        }

        Ok(())
    }

    fn versions<'a>(&'a self) -> Box<dyn Iterator<Item = Version> + 'a> {
        Box::new(self.versions.keys().cloned())
    }

    fn chain_info(&mut self) -> ChainInfo<'_> {
        ChainInfo::new(
            &mut self.root_account,
            &mut self.treasury_compliance_account,
            &mut self.designated_dealer_account,
            self.validators
                .values()
                .map(|v| v.json_rpc_endpoint().to_string())
                .next()
                .unwrap(),
            self.chain_id,
        )
    }

    fn logs_location(&mut self) -> String {
        self.dir.persist();
        self.dir.display().to_string()
    }
}
<commit_msg>forge: add some helpers for LocalSwarm<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    ChainInfo, FullNode, HealthCheckError, LocalNode, LocalVersion, Node, Swarm, SwarmExt,
    Validator, Version,
};
use anyhow::{anyhow, Result};
use diem_config::config::NodeConfig;
use diem_genesis_tool::{fullnode_builder::FullnodeConfig, validator_builder::ValidatorBuilder};
use diem_sdk::{
    crypto::ed25519::Ed25519PrivateKey,
    types::{
        chain_id::ChainId, transaction::Transaction, waypoint::Waypoint, AccountKey, LocalAccount,
        PeerId,
    },
};
use std::{
    collections::HashMap,
    convert::TryFrom,
    fs, mem,
    num::NonZeroUsize,
    ops,
    path::{Path, PathBuf},
    sync::Arc,
    time::{Duration, Instant},
};
use tempfile::TempDir;

#[derive(Debug)]
pub enum SwarmDirectory {
    Persistent(PathBuf),
    Temporary(TempDir),
}

impl SwarmDirectory {
    pub fn persist(&mut self) {
        match self {
            SwarmDirectory::Persistent(_) => {}
            SwarmDirectory::Temporary(_) => {
                let mut temp = SwarmDirectory::Persistent(PathBuf::new());
                mem::swap(self, &mut temp);
                let _ = mem::replace(self, temp.into_persistent());
            }
        }
    }

    pub fn into_persistent(self) -> Self {
        match self {
            SwarmDirectory::Temporary(tempdir) => SwarmDirectory::Persistent(tempdir.into_path()),
            SwarmDirectory::Persistent(dir) => SwarmDirectory::Persistent(dir),
        }
    }
}

impl ops::Deref for SwarmDirectory {
    type Target = Path;

    fn deref(&self) -> &Self::Target {
        match self {
            SwarmDirectory::Persistent(dir) => dir.deref(),
            SwarmDirectory::Temporary(dir) => dir.path(),
        }
    }
}

impl AsRef<Path> for SwarmDirectory {
    fn as_ref(&self) -> &Path {
        match self {
            SwarmDirectory::Persistent(dir) => dir.as_ref(),
            SwarmDirectory::Temporary(dir) => dir.as_ref(),
        }
    }
}

#[derive(Debug)]
struct ValidatorNetworkAddressEncryptionKey {
    key: diem_sdk::types::network_address::encrypted::Key,
    version: diem_sdk::types::network_address::encrypted::KeyVersion,
}

pub struct LocalSwarmBuilder {
    versions: Arc<HashMap<Version, LocalVersion>>,
    initial_version: Option<Version>,
    template: NodeConfig,
    number_of_validators: NonZeroUsize,
    dir: Option<PathBuf>,
}

impl LocalSwarmBuilder {
    pub fn new(versions: Arc<HashMap<Version, LocalVersion>>) -> Self {
        Self {
            versions,
            initial_version: None,
            template: NodeConfig::default_for_validator(),
            number_of_validators: NonZeroUsize::new(1).unwrap(),
            dir: None,
        }
    }

    pub fn initial_version(mut self, initial_version: Version) -> Self {
        self.initial_version = Some(initial_version);
        self
    }

    pub fn template(mut self, template: NodeConfig) -> Self {
        self.template = template;
        self
    }

    pub fn number_of_validators(mut self, number_of_validators: NonZeroUsize) -> Self {
        self.number_of_validators = number_of_validators;
        self
    }

    pub fn dir<T: AsRef<Path>>(mut self, dir: T) -> Self {
        self.dir = Some(dir.as_ref().into());
        self
    }

    pub fn build<R>(self, rng: R) -> Result<LocalSwarm>
    where
        R: ::rand::RngCore + ::rand::CryptoRng,
    {
        let dir = if let Some(dir) = self.dir {
            if dir.exists() {
                fs::remove_dir_all(&dir)?;
            }
            fs::create_dir_all(&dir)?;
            SwarmDirectory::Persistent(dir)
        } else {
            SwarmDirectory::Temporary(TempDir::new()?)
        };

        let (root_keys, genesis, genesis_waypoint, validators) = ValidatorBuilder::new(
            &dir,
            diem_framework_releases::current_module_blobs().to_vec(),
        )
        .num_validators(self.number_of_validators)
        .template(self.template)
        .build(rng)?;

        // Get the initial version to start the nodes with, either the one provided or fallback to
        // using the the latest version
        let versions = self.versions;
        let initial_version = self.initial_version.unwrap_or_else(|| {
            versions
                .iter()
                .max_by(|v1, v2| v1.0.cmp(v2.0))
                .unwrap()
                .0
                .clone()
        });
        let version = versions.get(&initial_version).unwrap();

        let validators = validators
            .into_iter()
            .map(|v| {
                let node = LocalNode::new(version.to_owned(), v.name, v.directory)?;
                Ok((node.peer_id(), node))
            })
            .collect::<Result<HashMap<_, _>>>()?;

        let validator_network_address_encryption_key = ValidatorNetworkAddressEncryptionKey {
            key: root_keys.validator_network_address_encryption_key,
            version: root_keys.validator_network_address_encryption_key_version,
        };

        let root_account = LocalAccount::new(
            diem_sdk::types::account_config::diem_root_address(),
            AccountKey::from_private_key(root_keys.root_key),
            1,
        );

        // Designated dealer account is built to have the same private key as the TC account
        let designated_dealer_account = LocalAccount::new(
            diem_sdk::types::account_config::testnet_dd_account_address(),
            AccountKey::from_private_key(
                Ed25519PrivateKey::try_from(root_keys.treasury_compliance_key.to_bytes().as_ref())
                    .unwrap(),
            ),
            0,
        );

        let treasury_compliance_account = LocalAccount::new(
            diem_sdk::types::account_config::treasury_compliance_account_address(),
            AccountKey::from_private_key(root_keys.treasury_compliance_key),
            0,
        );

        Ok(LocalSwarm {
            node_name_counter: validators.len() as u64,
            genesis,
            genesis_waypoint,
            versions,
            validators,
            fullnodes: HashMap::new(),
            dir,
            validator_network_address_encryption_key,
            root_account,
            treasury_compliance_account,
            designated_dealer_account,
            chain_id: ChainId::test(),
        })
    }
}

#[derive(Debug)]
pub struct LocalSwarm {
    node_name_counter: u64,
    genesis: Transaction,
    genesis_waypoint: Waypoint,
    versions: Arc<HashMap<Version, LocalVersion>>,
    validators: HashMap<PeerId, LocalNode>,
    fullnodes: HashMap<PeerId, LocalNode>,
    dir: SwarmDirectory,
    validator_network_address_encryption_key: ValidatorNetworkAddressEncryptionKey,
    root_account: LocalAccount,
    treasury_compliance_account: LocalAccount,
    designated_dealer_account: LocalAccount,
    chain_id: ChainId,
}

impl LocalSwarm {
    pub fn builder(versions: Arc<HashMap<Version, LocalVersion>>) -> LocalSwarmBuilder {
        LocalSwarmBuilder::new(versions)
    }

    pub fn launch(&mut self) -> Result<()> {
        // Start all the validators
        for validator in self.validators.values_mut() {
            validator.start()?;
        }

        // Wait for all of them to startup
        let deadline = Instant::now() + Duration::from_secs(60);
        self.wait_for_startup()?;
        self.wait_for_connectivity(deadline)?;
        self.liveness_check(deadline)?;

        Ok(())
    }

    fn wait_for_startup(&mut self) -> Result<()> {
        let num_attempts = 10;
        let mut done = vec![false; self.validators.len()];
        for i in 0..num_attempts {
            println!("Wait for startup attempt: {} of {}", i, num_attempts);
            for (node, done) in self.validators.values_mut().zip(done.iter_mut()) {
                if *done {
                    continue;
                }
                match node.health_check() {
                    Ok(()) => *done = true,

                    Err(HealthCheckError::Unknown(e)) => {
                        return Err(anyhow!("Diem node '{}' is not running: {}", node.name(), e));
                    }
                    Err(HealthCheckError::NotRunning) => {
                        return Err(anyhow!("Diem node '{}' is not running", node.name()));
                    }
                    Err(HealthCheckError::RpcFailure(e)) => {
                        println!("rpc error: {}", e);
                        continue;
                    }
                }
            }

            // Check if all the nodes have been successfully launched
            if done.iter().all(|status| *status) {
                return Ok(());
            }

            ::std::thread::sleep(::std::time::Duration::from_millis(1000));
        }

        Err(anyhow!("Launching Swarm timed out"))
    }

    fn add_fullnode(&mut self, version: &Version, template: NodeConfig) -> Result<PeerId> {
        let name = self.node_name_counter.to_string();
        self.node_name_counter += 1;
        let fullnode_config = FullnodeConfig::public_fullnode(
            name,
            self.dir.as_ref(),
            template,
            &self.genesis_waypoint,
            &self.genesis,
        )?;

        let version = self.versions.get(version).unwrap();
        let mut fullnode = LocalNode::new(
            version.to_owned(),
            fullnode_config.name,
            fullnode_config.directory,
        )?;

        let peer_id = fullnode.peer_id();
        fullnode.start()?;

        self.fullnodes.insert(peer_id, fullnode);

        Ok(peer_id)
    }

    pub fn chain_id(&self) -> ChainId {
        self.chain_id
    }

    pub fn validators(&self) -> impl Iterator<Item = &LocalNode> {
        self.validators.values()
    }

    pub fn validators_mut(&mut self) -> impl Iterator<Item = &mut LocalNode> {
        self.validators.values_mut()
    }

    pub fn dir(&self) -> &Path {
        self.dir.as_ref()
    }
}

impl Swarm for LocalSwarm {
    fn health_check(&mut self) -> Result<()> {
        todo!()
    }

    fn validators<'a>(&'a self) -> Box<dyn Iterator<Item = &'a dyn Validator> + 'a> {
        Box::new(self.validators.values().map(|v| v as &'a dyn Validator))
    }

    fn validators_mut<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a mut dyn Validator> + 'a> {
        Box::new(
            self.validators
                .values_mut()
                .map(|v| v as &'a mut dyn Validator),
        )
    }

    fn validator(&self, id: PeerId) -> Option<&dyn Validator> {
        self.validators.get(&id).map(|v| v as &dyn Validator)
    }

    fn validator_mut(&mut self, id: PeerId) -> Option<&mut dyn Validator> {
        self.validators
            .get_mut(&id)
            .map(|v| v as &mut dyn Validator)
    }

    fn upgrade_validator(&mut self, id: PeerId, version: &Version) -> Result<()> {
        let version = self
            .versions
            .get(version)
            .cloned()
            .ok_or_else(|| anyhow!("Invalid version: {:?}", version))?;
        let validator = self
            .validators
            .get_mut(&id)
            .ok_or_else(|| anyhow!("Invalid id: {}", id))?;
        validator.upgrade(version)
    }

    fn full_nodes<'a>(&'a self) -> Box<dyn Iterator<Item = &'a dyn FullNode> + 'a> {
        Box::new(self.fullnodes.values().map(|v| v as &'a dyn FullNode))
    }

    fn full_nodes_mut<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a mut dyn FullNode> + 'a> {
        Box::new(
            self.fullnodes
                .values_mut()
                .map(|v| v as &'a mut dyn FullNode),
        )
    }

    fn full_node(&self, id: PeerId) -> Option<&dyn FullNode> {
        self.fullnodes.get(&id).map(|v| v as &dyn FullNode)
    }

    fn full_node_mut(&mut self, id: PeerId) -> Option<&mut dyn FullNode> {
        self.fullnodes.get_mut(&id).map(|v| v as &mut dyn FullNode)
    }

    fn add_validator(&mut self, _version: &Version, _template: NodeConfig) -> Result<PeerId> {
        todo!()
    }

    fn remove_validator(&mut self, _id: PeerId) -> Result<()> {
        todo!()
    }

    fn add_full_node(&mut self, version: &Version, template: NodeConfig) -> Result<PeerId> {
        self.add_fullnode(version, template)
    }

    fn remove_full_node(&mut self, id: PeerId) -> Result<()> {
        if let Some(mut fullnode) = self.fullnodes.remove(&id) {
            fullnode.stop();
        }

        Ok(())
    }

    fn versions<'a>(&'a self) -> Box<dyn Iterator<Item = Version> + 'a> {
        Box::new(self.versions.keys().cloned())
    }

    fn chain_info(&mut self) -> ChainInfo<'_> {
        ChainInfo::new(
            &mut self.root_account,
            &mut self.treasury_compliance_account,
            &mut self.designated_dealer_account,
            self.validators
                .values()
                .map(|v| v.json_rpc_endpoint().to_string())
                .next()
                .unwrap(),
            self.chain_id,
        )
    }

    fn logs_location(&mut self) -> String {
        self.dir.persist();
        self.dir.display().to_string()
    }
}
<|endoftext|>


<commit_before>use orbclient::Renderer;
use std::cell::{Cell, RefCell};
use std::cmp::max;
use std::sync::Arc;

use cell::{CloneCell, CheckSet};
use draw::draw_box;
use event::Event;
use point::Point;
use rect::Rect;
use theme::{Theme, Selector};
use traits::{Click, Place, Text};
use widgets::Widget;

pub struct Menu {
    pub rect: Cell<Rect>,
    text: CloneCell<String>,
    text_offset: Cell<Point>,
    entries: RefCell<Vec<Arc<Entry>>>,
    click_callback: RefCell<Option<Arc<Fn(&Menu, Point)>>>,
    pressed: Cell<bool>,
    activated: Cell<bool>,
}

pub struct Separator {
    pub rect: Cell<Rect>,
}

pub trait Entry: Widget {
    fn entry_text(&self) -> String;
}

impl Menu {
    pub fn new<S: Into<String>>(name: S) -> Arc<Self> {
        Arc::new(Menu {
            rect: Cell::new(Rect::default()),
            text: CloneCell::new(name.into()),
            text_offset: Cell::new(Point::default()),
            entries: RefCell::new(Vec::new()),
            click_callback: RefCell::new(None),
            pressed: Cell::new(false),
            activated: Cell::new(false),
        })
    }

    pub fn add<T: Entry>(&self, new_entry: &Arc<T>) {
        let mut rect = self.rect.get();
        let text_width = new_entry.entry_text().len() as u32 * 8;
        if rect.width < text_width {
            rect.width = text_width;
        }

        let mut y = rect.y + rect.height as i32;
        for entry in self.entries.borrow().iter() {
            let mut entry_rect = entry.rect().get();
            y += entry_rect.height as i32;

            if entry_rect.width < rect.width {
                entry_rect.width = rect.width;
                entry.rect().set(entry_rect);
            } else {
                rect.width = entry_rect.width;
            }
        }
        rect.y = y;
        new_entry.rect().set(rect);
        self.entries.borrow_mut().push(new_entry.clone());
    }
}

impl Text for Menu {
    fn text<S: Into<String>>(&self, text: S) -> &Self {
        self.text.set(text.into());
        self
    }

    fn text_offset(&self, x: i32, y: i32) -> &Self {
        self.text_offset.set(Point::new(x, y));
        self
    }
}

impl Click for Menu {
    fn emit_click(&self, point: Point) {
        if let Some(ref click_callback) = *self.click_callback.borrow() {
            click_callback(self, point);
        }
    }

    fn on_click<T: Fn(&Self, Point) + 'static>(&self, func: T) -> &Self {
        *self.click_callback.borrow_mut() = Some(Arc::new(func));
        self
    }
}

impl Place for Menu {}

impl Widget for Menu {
    fn name(&self) -> &str {
        "Menu"
    }

    fn rect(&self) -> &Cell<Rect> {
        &self.rect
    }

    fn draw(&self, renderer: &mut Renderer, _focused: bool, theme: &Theme) {
        let rect = self.rect.get();

        if self.activated.get() {
            draw_box(renderer, rect, theme, &Selector::new(Some("menu-button")).with_pseudo_class("active"));
        } else {
            draw_box(renderer, rect, theme, &Selector::new(Some("menu-button")).with_pseudo_class("inactive"));
        }

        let text = self.text.borrow();
        let mut point = self.text_offset.get();
        for c in text.chars() {
            if c == '\n' {
                point.x = self.text_offset.get().x;
                point.y += 16;
            } else {
                if point.x + 8 <= rect.width as i32 && point.y + 16 <= rect.height as i32 {
                    renderer.char(point.x + rect.x, point.y + rect.y, c, theme.color("color", &"button".into()));
                }
                point.x += 8;
            }
        }

        if self.activated.get() {
            let mut max_width = 0;
            let mut max_height = 0;

            for entry in self.entries.borrow().iter() {
                let r = entry.rect().get();
                max_width = max(max_width, r.x + r.width as i32 - rect.x);
                max_height = max(max_height, r.y + r.height as i32 - rect.y - rect.height as i32);
            }

            renderer.rect(rect.x - 1, rect.y + rect.height as i32 - 1, max_width as u32 + 2, max_height as u32 + 2, self.fg_border);

            for entry in self.entries.borrow().iter() {
                entry.draw(renderer, _focused, theme);
            }
        }
    }

    fn event(&self, event: Event, focused: bool, redraw: &mut bool) -> bool {
        let mut ignore_event = false;
        if self.activated.get() {
            for entry in self.entries.borrow().iter() {
                if entry.event(event, focused, redraw) {
                    ignore_event = true;
                    self.pressed.set(true);
                }
            }
        }

        match event {
            Event::Mouse { point, left_button, .. } => {
                let mut click = false;

                let rect = self.rect.get();
                if rect.contains(point) {
                    if left_button {
                        self.pressed.set(!self.pressed.get());

                        if self.activated.check_set(true) {
                            click = true;
                            *redraw = true;
                        }
                    } else {
                        if !self.pressed.get() {
                            if self.activated.check_set(false) {
                                click = true;
                                *redraw = true;
                            }
                        }
                    }
                } else {
                    if !ignore_event {
                        if left_button {
                            self.pressed.set(false);
                        } else {
                            if !self.pressed.get() {
                                if self.activated.check_set(false) {
                                    *redraw = true;
                                }
                            }
                        }
                    }
                }

                if click {
                    let click_point: Point = point - rect.point();
                    self.emit_click(click_point);
                }
            }
            _ => (),
        }
        focused
    }
}

pub struct Action {
    rect: Cell<Rect>,
    text: CloneCell<String>,
    text_offset: Cell<Point>,
    click_callback: RefCell<Option<Arc<Fn(&Action, Point)>>>,
    pressed: Cell<bool>,
    hover: Cell<bool>,
}

impl Action {
    pub fn new<S: Into<String>>(text: S) -> Arc<Self> {
        Arc::new(Action {
            rect: Cell::new(Rect::default()),
            text: CloneCell::new(text.into()),
            text_offset: Cell::new(Point::default()),
            click_callback: RefCell::new(None),
            pressed: Cell::new(false),
            hover: Cell::new(false),
        })
    }
}

impl Click for Action {
    fn emit_click(&self, point: Point) {
        if let Some(ref click_callback) = *self.click_callback.borrow() {
            click_callback(self, point);
        }
    }

    fn on_click<T: Fn(&Self, Point) + 'static>(&self, func: T) -> &Self {
        *self.click_callback.borrow_mut() = Some(Arc::new(func));
        self
    }
}

impl Text for Action {
    fn text<S: Into<String>>(&self, text: S) -> &Self {
        self.text.set(text.into());
        self
    }

    fn text_offset(&self, x: i32, y: i32) -> &Self {
        self.text_offset.set(Point::new(x, y));
        self
    }
}

impl Widget for Action {
    fn name(&self) -> &str {
        "Action"
    }

    fn rect(&self) -> &Cell<Rect> {
        &self.rect
    }

    fn draw(&self, renderer: &mut Renderer, _focused: bool, theme: &Theme) {
        let rect = self.rect.get();

        let pseudo_class = if self.hover.get() { "active" } else { "inactive" };

        draw_box(renderer, rect, theme, &Selector::new(Some("action")).with_pseudo_class(pseudo_class));

        let text = self.text.borrow();
        let mut point = self.text_offset.get();
        for c in text.chars() {
            if c == '\n' {
                point.x = self.text_offset.get().x;
                point.y += 16;
            } else {
                if point.x + 8 <= rect.width as i32 && point.y + 16 <= rect.height as i32 {
                    renderer.char(point.x + rect.x, point.y + rect.y, c, theme.color("color", &"action".into()));
                }
                point.x += 8;
            }
        }
    }

    fn event(&self, event: Event, _focused: bool, redraw: &mut bool) -> bool {
        match event {
            Event::Mouse { point, left_button, .. } => {
                let mut click = false;
                let rect = self.rect.get();

                if rect.contains(point) {
                    if self.hover.check_set(true) {
                        *redraw = true;
                    }

                    if left_button {
                        if self.pressed.check_set(true) {
                            *redraw = true;
                        }
                    } else {
                        if self.pressed.check_set(false) {
                            click = true;
                            self.hover.set(false);
                            *redraw = true;
                        }
                    }
                } else {
                    if self.hover.check_set(false) {
                        *redraw = true;
                    }

                    if !left_button {
                        if self.pressed.check_set(false) {
                            *redraw = true;
                        }
                    }
                }

                if click {
                    let click_point: Point = point - rect.point();
                    self.emit_click(click_point);
                }
            }
            _ => (),
        }

        false
    }
}

impl Entry for Action {
    fn entry_text(&self) -> String {
        self.text.get()
    }
}

impl Separator {
    pub fn new() -> Arc<Self> {
        Arc::new(Separator {
            rect: Cell::new(Rect::default()),
        })
    }
}

impl Widget for Separator {
    fn name(&self) -> &str {
        "Separator"
    }

    fn rect(&self) -> &Cell<Rect> {
        &self.rect
    }

    fn draw(&self, renderer: &mut Renderer, _focused: bool, theme: &Theme) {
        let rect = self.rect.get();
        draw_box(renderer, rect, theme, &"separator".into());

        let line_y = rect.y + rect.height as i32 / 2;
        renderer.rect(rect.x, line_y, rect.width, 1, theme.color("color", &"separator".into()));
    }

    fn event(&self, event: Event, _focused: bool, _redraw: &mut bool) -> bool {
        let mut ignore_event = false;
        match event {
            Event::Mouse { point, .. } => {
                let rect = self.rect.get();
                if rect.contains(point) {
                    ignore_event = true;
                }
            }
            _ => (),
        }
        ignore_event
    }
}

impl Entry for Separator {
    fn entry_text(&self) -> String {
        String::new()
    }
}
<commit_msg>Fix menu widget<commit_after>use orbclient::Renderer;
use std::cell::{Cell, RefCell};
use std::cmp::max;
use std::sync::Arc;

use cell::{CloneCell, CheckSet};
use draw::draw_box;
use event::Event;
use point::Point;
use rect::Rect;
use theme::{Theme, Selector};
use traits::{Click, Place, Text};
use widgets::Widget;

pub struct Menu {
    pub rect: Cell<Rect>,
    text: CloneCell<String>,
    text_offset: Cell<Point>,
    entries: RefCell<Vec<Arc<Entry>>>,
    click_callback: RefCell<Option<Arc<Fn(&Menu, Point)>>>,
    pressed: Cell<bool>,
    activated: Cell<bool>,
}

pub struct Separator {
    pub rect: Cell<Rect>,
}

pub trait Entry: Widget {
    fn entry_text(&self) -> String;
}

impl Menu {
    pub fn new<S: Into<String>>(name: S) -> Arc<Self> {
        Arc::new(Menu {
            rect: Cell::new(Rect::default()),
            text: CloneCell::new(name.into()),
            text_offset: Cell::new(Point::default()),
            entries: RefCell::new(Vec::new()),
            click_callback: RefCell::new(None),
            pressed: Cell::new(false),
            activated: Cell::new(false),
        })
    }

    pub fn add<T: Entry>(&self, new_entry: &Arc<T>) {
        let mut rect = self.rect.get();
        let text_width = new_entry.entry_text().len() as u32 * 8;
        if rect.width < text_width {
            rect.width = text_width;
        }

        let mut y = rect.y + rect.height as i32;
        for entry in self.entries.borrow().iter() {
            let mut entry_rect = entry.rect().get();
            y += entry_rect.height as i32;

            if entry_rect.width < rect.width {
                entry_rect.width = rect.width;
                entry.rect().set(entry_rect);
            } else {
                rect.width = entry_rect.width;
            }
        }
        rect.y = y;
        new_entry.rect().set(rect);
        self.entries.borrow_mut().push(new_entry.clone());
    }
}

impl Text for Menu {
    fn text<S: Into<String>>(&self, text: S) -> &Self {
        self.text.set(text.into());
        self
    }

    fn text_offset(&self, x: i32, y: i32) -> &Self {
        self.text_offset.set(Point::new(x, y));
        self
    }
}

impl Click for Menu {
    fn emit_click(&self, point: Point) {
        if let Some(ref click_callback) = *self.click_callback.borrow() {
            click_callback(self, point);
        }
    }

    fn on_click<T: Fn(&Self, Point) + 'static>(&self, func: T) -> &Self {
        *self.click_callback.borrow_mut() = Some(Arc::new(func));
        self
    }
}

impl Place for Menu {}

impl Widget for Menu {
    fn name(&self) -> &str {
        "Menu"
    }

    fn rect(&self) -> &Cell<Rect> {
        &self.rect
    }

    fn draw(&self, renderer: &mut Renderer, _focused: bool, theme: &Theme) {
        let rect = self.rect.get();

        if self.activated.get() {
            draw_box(renderer, rect, theme, &Selector::new(Some("menu-button")).with_pseudo_class("active"));

            let mut max_width = 0;
            let mut max_height = 0;

            for entry in self.entries.borrow().iter() {
                let r = entry.rect().get();
                max_width = max(max_width, r.x + r.width as i32 - rect.x);
                max_height = max(max_height, r.y + r.height as i32 - rect.y - rect.height as i32);
            }

            let entries_rect = Rect::new(
                rect.x - 1, rect.y + rect.height as i32 - 1,
                max_width as u32 + 2, max_height as u32 + 2,
            );

            draw_box(renderer, entries_rect, theme, &Selector::new(Some("menu-button")).with_pseudo_class("active"));
        } else {
            draw_box(renderer, rect, theme, &Selector::new(Some("menu-button")).with_pseudo_class("inactive"));
        }

        let text = self.text.borrow();
        let mut point = self.text_offset.get();
        for c in text.chars() {
            if c == '\n' {
                point.x = self.text_offset.get().x;
                point.y += 16;
            } else {
                if point.x + 8 <= rect.width as i32 && point.y + 16 <= rect.height as i32 {
                    renderer.char(point.x + rect.x, point.y + rect.y, c, theme.color("color", &"button".into()));
                }
                point.x += 8;
            }
        }

        if self.activated.get() {
            for entry in self.entries.borrow().iter() {
                entry.draw(renderer, _focused, theme);
            }
        }
    }

    fn event(&self, event: Event, focused: bool, redraw: &mut bool) -> bool {
        let mut ignore_event = false;
        if self.activated.get() {
            for entry in self.entries.borrow().iter() {
                if entry.event(event, focused, redraw) {
                    ignore_event = true;
                    self.pressed.set(true);
                }
            }
        }

        match event {
            Event::Mouse { point, left_button, .. } => {
                let mut click = false;

                let rect = self.rect.get();
                if rect.contains(point) {
                    if left_button {
                        self.pressed.set(!self.pressed.get());

                        if self.activated.check_set(true) {
                            click = true;
                            *redraw = true;
                        }
                    } else {
                        if !self.pressed.get() {
                            if self.activated.check_set(false) {
                                click = true;
                                *redraw = true;
                            }
                        }
                    }
                } else {
                    if !ignore_event {
                        if left_button {
                            self.pressed.set(false);
                        } else {
                            if !self.pressed.get() {
                                if self.activated.check_set(false) {
                                    *redraw = true;
                                }
                            }
                        }
                    }
                }

                if click {
                    let click_point: Point = point - rect.point();
                    self.emit_click(click_point);
                }
            }
            _ => (),
        }
        focused
    }
}

pub struct Action {
    rect: Cell<Rect>,
    text: CloneCell<String>,
    text_offset: Cell<Point>,
    click_callback: RefCell<Option<Arc<Fn(&Action, Point)>>>,
    pressed: Cell<bool>,
    hover: Cell<bool>,
}

impl Action {
    pub fn new<S: Into<String>>(text: S) -> Arc<Self> {
        Arc::new(Action {
            rect: Cell::new(Rect::default()),
            text: CloneCell::new(text.into()),
            text_offset: Cell::new(Point::default()),
            click_callback: RefCell::new(None),
            pressed: Cell::new(false),
            hover: Cell::new(false),
        })
    }
}

impl Click for Action {
    fn emit_click(&self, point: Point) {
        if let Some(ref click_callback) = *self.click_callback.borrow() {
            click_callback(self, point);
        }
    }

    fn on_click<T: Fn(&Self, Point) + 'static>(&self, func: T) -> &Self {
        *self.click_callback.borrow_mut() = Some(Arc::new(func));
        self
    }
}

impl Text for Action {
    fn text<S: Into<String>>(&self, text: S) -> &Self {
        self.text.set(text.into());
        self
    }

    fn text_offset(&self, x: i32, y: i32) -> &Self {
        self.text_offset.set(Point::new(x, y));
        self
    }
}

impl Widget for Action {
    fn name(&self) -> &str {
        "Action"
    }

    fn rect(&self) -> &Cell<Rect> {
        &self.rect
    }

    fn draw(&self, renderer: &mut Renderer, _focused: bool, theme: &Theme) {
        let rect = self.rect.get();

        let pseudo_class = if self.hover.get() { "active" } else { "inactive" };

        draw_box(renderer, rect, theme, &Selector::new(Some("action")).with_pseudo_class(pseudo_class));

        let text = self.text.borrow();
        let mut point = self.text_offset.get();
        for c in text.chars() {
            if c == '\n' {
                point.x = self.text_offset.get().x;
                point.y += 16;
            } else {
                if point.x + 8 <= rect.width as i32 && point.y + 16 <= rect.height as i32 {
                    renderer.char(point.x + rect.x, point.y + rect.y, c, theme.color("color", &"action".into()));
                }
                point.x += 8;
            }
        }
    }

    fn event(&self, event: Event, _focused: bool, redraw: &mut bool) -> bool {
        match event {
            Event::Mouse { point, left_button, .. } => {
                let mut click = false;
                let rect = self.rect.get();

                if rect.contains(point) {
                    if self.hover.check_set(true) {
                        *redraw = true;
                    }

                    if left_button {
                        if self.pressed.check_set(true) {
                            *redraw = true;
                        }
                    } else {
                        if self.pressed.check_set(false) {
                            click = true;
                            self.hover.set(false);
                            *redraw = true;
                        }
                    }
                } else {
                    if self.hover.check_set(false) {
                        *redraw = true;
                    }

                    if !left_button {
                        if self.pressed.check_set(false) {
                            *redraw = true;
                        }
                    }
                }

                if click {
                    let click_point: Point = point - rect.point();
                    self.emit_click(click_point);
                }
            }
            _ => (),
        }

        false
    }
}

impl Entry for Action {
    fn entry_text(&self) -> String {
        self.text.get()
    }
}

impl Separator {
    pub fn new() -> Arc<Self> {
        Arc::new(Separator {
            rect: Cell::new(Rect::default()),
        })
    }
}

impl Widget for Separator {
    fn name(&self) -> &str {
        "Separator"
    }

    fn rect(&self) -> &Cell<Rect> {
        &self.rect
    }

    fn draw(&self, renderer: &mut Renderer, _focused: bool, theme: &Theme) {
        let rect = self.rect.get();
        draw_box(renderer, rect, theme, &"separator".into());

        let line_y = rect.y + rect.height as i32 / 2;
        renderer.rect(rect.x, line_y, rect.width, 1, theme.color("color", &"separator".into()));
    }

    fn event(&self, event: Event, _focused: bool, _redraw: &mut bool) -> bool {
        let mut ignore_event = false;
        match event {
            Event::Mouse { point, .. } => {
                let rect = self.rect.get();
                if rect.contains(point) {
                    ignore_event = true;
                }
            }
            _ => (),
        }
        ignore_event
    }
}

impl Entry for Separator {
    fn entry_text(&self) -> String {
        String::new()
    }
}
<|endoftext|>


<commit_before>use std::iter::FromIterator;

use serde::{
    Deserialize,
    Serialize,
};

use thiserror::Error;

use rnc_core::{
    urs,
    urs_taxid,
};

use crate::normalize::{
    accession::{
        AccessionVec,
        CrossReference,
        RawAccession,
        ReferenceVec,
    },
    basic::Basic,
    crs::{
        Crs,
        CrsVec,
    },
    feedback::{
        Feedback,
        FeedbackVec,
    },
    go_annotation::GoAnnotation,
    interacting_protein::InteractingProtein,
    interacting_rna::InteractingRna,
    precompute::{
        Precompute,
        PrecomputeSummary,
    },
    qa_status::QaStatus,
    r2dt::R2dt,
    rfam_hit::{
        RfamHit,
        RfamHitVec,
    },
    so_tree,
};

#[derive(Error, Debug)]
pub enum NormalizationError {
    #[error("Could not parse {0}")]
    UrsParsingError(#[from] urs::Error),

    #[error("Could not parse {0}")]
    UrsTaxidParsingError(#[from] urs_taxid::Error),
}

#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Raw {
    pub id: usize,
    pub base: Basic,
    pub crs: Vec<Crs>,
    pub feedback: Vec<Feedback>,
    pub go_annotations: Vec<GoAnnotation>,
    pub interacting_proteins: Vec<InteractingProtein>,
    pub interacting_rnas: Vec<InteractingRna>,
    pub precompute: Precompute,
    pub qa_status: QaStatus,
    pub r2dt: Option<R2dt>,
    pub rfam_hits: Vec<RfamHit>,
    pub so_tree: so_tree::SoTree,
}

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct Normalized {
    urs: String,
    taxid: usize,
    urs_taxid: String,
    short_urs: String,
    deleted: String,

    so_rna_type_tree: so_tree::SoTree,

    #[serde(flatten)]
    pre_summary: PrecomputeSummary,

    #[serde(flatten)]
    basic: Basic,

    // #[serde(flatten)]
    // dates: Dates,
    qa_status: QaStatus,

    secondary: Option<R2dt>,

    #[serde(flatten)]
    accessions: AccessionVec,
    cross_references: Vec<CrossReference>,
    crs: CrsVec,
    overlaps: FeedbackVec,
    go_annotations: Vec<GoAnnotation>,

    interacting_proteins: Vec<InteractingProtein>,

    interacting_rnas: Vec<InteractingRna>,

    #[serde(flatten)]
    references: ReferenceVec,

    #[serde(flatten)]
    rfam_hits: RfamHitVec,
}

impl Raw {
    pub fn urs_taxid(&self) -> String {
        return self.base.urs_taxid.to_owned();
    }

    pub fn urs(&self) -> Result<String, urs_taxid::Error> {
        let ut: urs_taxid::UrsTaxid = self.urs_taxid().parse()?;
        let urs: urs::Urs = ut.into();
        Ok(urs.to_string())
    }

    pub fn taxid(&self) -> Result<u64, urs_taxid::Error> {
        let ut: urs_taxid::UrsTaxid = self.urs_taxid().parse()?;
        Ok(ut.taxid())
    }

    pub fn short_urs(&self) -> Result<String, urs_taxid::Error> {
        let ut: urs_taxid::UrsTaxid = self.urs_taxid().parse()?;
        let urs: urs::Urs = ut.into();
        Ok(urs.short_urs())
    }

    pub fn short_urs_taxid(&self) -> Result<String, urs_taxid::Error> {
        let ut: urs_taxid::UrsTaxid = self.urs_taxid().parse()?;
        Ok(ut.short())
    }
}

impl Normalized {
    pub fn new(raw: Raw, accessions: Vec<RawAccession>) -> Result<Self, NormalizationError> {
        let pre_summary = PrecomputeSummary::from(raw.precompute);
        let base = raw.base.clone();

        let urs_taxid = base.urs_taxid.to_owned();
        assert!(
            urs_taxid == accessions[0].urs_taxid,
            "Accession urs {} and base urs {} disagree for id {}",
            accessions[0].urs_taxid,
            urs_taxid,
            base.id
        );
        let parts: Vec<&str> = urs_taxid.split("_").collect();
        let urs = parts[0].to_owned();
        let taxid = parts[1].parse::<usize>().unwrap();
        let parsed: urs_taxid::UrsTaxid = urs_taxid.parse()?;
        let references = ReferenceVec::from_iter(accessions.clone());

        Ok(Self {
            urs_taxid,
            urs,
            taxid,
            short_urs: parsed.short(),
            deleted: String::from("N"),

            so_rna_type_tree: raw.so_tree,

            pre_summary,
            basic: base,
            qa_status: raw.qa_status,
            secondary: raw.r2dt,

            accessions: AccessionVec::from_iter(accessions.clone()),
            cross_references: accessions.into_iter().map(CrossReference::from).collect(),
            crs: CrsVec::from_iter(raw.crs.clone()),
            overlaps: FeedbackVec::from_iter(raw.feedback.clone()),
            go_annotations: raw.go_annotations.clone(),
            interacting_proteins: raw.interacting_proteins,
            interacting_rnas: raw.interacting_rnas,
            references,
            rfam_hits: RfamHitVec::from_iter(raw.rfam_hits.clone()),
        })
    }
}
<commit_msg>Use parsed urs data<commit_after>use std::iter::FromIterator;
use std::convert::TryInto;

use serde::{
    Deserialize,
    Serialize,
};

use thiserror::Error;

use rnc_core::{
    urs,
    urs_taxid,
};

use crate::normalize::{
    accession::{
        AccessionVec,
        CrossReference,
        RawAccession,
        ReferenceVec,
    },
    basic::Basic,
    crs::{
        Crs,
        CrsVec,
    },
    feedback::{
        Feedback,
        FeedbackVec,
    },
    go_annotation::GoAnnotation,
    interacting_protein::InteractingProtein,
    interacting_rna::InteractingRna,
    precompute::{
        Precompute,
        PrecomputeSummary,
    },
    qa_status::QaStatus,
    r2dt::R2dt,
    rfam_hit::{
        RfamHit,
        RfamHitVec,
    },
    so_tree,
};

#[derive(Error, Debug)]
pub enum NormalizationError {
    #[error("Could not parse {0}")]
    UrsParsingError(#[from] urs::Error),

    #[error("Could not parse {0}")]
    UrsTaxidParsingError(#[from] urs_taxid::Error),
}

#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Raw {
    pub id: usize,
    pub base: Basic,
    pub crs: Vec<Crs>,
    pub feedback: Vec<Feedback>,
    pub go_annotations: Vec<GoAnnotation>,
    pub interacting_proteins: Vec<InteractingProtein>,
    pub interacting_rnas: Vec<InteractingRna>,
    pub precompute: Precompute,
    pub qa_status: QaStatus,
    pub r2dt: Option<R2dt>,
    pub rfam_hits: Vec<RfamHit>,
    pub so_tree: so_tree::SoTree,
}

#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct Normalized {
    urs: String,
    taxid: usize,
    urs_taxid: String,
    short_urs: String,
    deleted: String,

    so_rna_type_tree: so_tree::SoTree,

    #[serde(flatten)]
    pre_summary: PrecomputeSummary,

    #[serde(flatten)]
    basic: Basic,

    // #[serde(flatten)]
    // dates: Dates,
    qa_status: QaStatus,

    secondary: Option<R2dt>,

    #[serde(flatten)]
    accessions: AccessionVec,
    cross_references: Vec<CrossReference>,
    crs: CrsVec,
    overlaps: FeedbackVec,
    go_annotations: Vec<GoAnnotation>,

    interacting_proteins: Vec<InteractingProtein>,

    interacting_rnas: Vec<InteractingRna>,

    #[serde(flatten)]
    references: ReferenceVec,

    #[serde(flatten)]
    rfam_hits: RfamHitVec,
}

impl Raw {
    pub fn urs_taxid(&self) -> String {
        return self.base.urs_taxid.to_owned();
    }

    pub fn urs(&self) -> Result<String, urs_taxid::Error> {
        let ut: urs_taxid::UrsTaxid = self.urs_taxid().parse()?;
        let urs: urs::Urs = ut.into();
        Ok(urs.to_string())
    }

    pub fn taxid(&self) -> Result<u64, urs_taxid::Error> {
        let ut: urs_taxid::UrsTaxid = self.urs_taxid().parse()?;
        Ok(ut.taxid())
    }

    pub fn short_urs(&self) -> Result<String, urs_taxid::Error> {
        let ut: urs_taxid::UrsTaxid = self.urs_taxid().parse()?;
        let urs: urs::Urs = ut.into();
        Ok(urs.short_urs())
    }

    pub fn short_urs_taxid(&self) -> Result<String, urs_taxid::Error> {
        let ut: urs_taxid::UrsTaxid = self.urs_taxid().parse()?;
        Ok(ut.short())
    }
}

impl Normalized {
    pub fn new(raw: Raw, accessions: Vec<RawAccession>) -> Result<Self, NormalizationError> {
        let pre_summary = PrecomputeSummary::from(raw.precompute);
        let base = raw.base.clone();

        let urs_taxid = base.urs_taxid.to_owned();
        assert!(
            urs_taxid == accessions[0].urs_taxid,
            "Accession urs {} and base urs {} disagree for id {}",
            accessions[0].urs_taxid,
            urs_taxid,
            base.id
        );
        let parsed: urs_taxid::UrsTaxid = urs_taxid.parse()?;
        let references = ReferenceVec::from_iter(accessions.clone());

        Ok(Self {
            urs_taxid,
            urs: parsed.urs().to_string(),
            taxid: parsed.taxid().try_into().unwrap(),
            short_urs: parsed.short(),
            deleted: String::from("N"),
            so_rna_type_tree: raw.so_tree,
            pre_summary,
            basic: base,
            qa_status: raw.qa_status,
            secondary: raw.r2dt,
            accessions: AccessionVec::from_iter(accessions.clone()),
            cross_references: accessions.into_iter().map(CrossReference::from).collect(),
            crs: CrsVec::from_iter(raw.crs.clone()),
            overlaps: FeedbackVec::from_iter(raw.feedback.clone()),
            go_annotations: raw.go_annotations.clone(),
            interacting_proteins: raw.interacting_proteins,
            interacting_rnas: raw.interacting_rnas,
            references,
            rfam_hits: RfamHitVec::from_iter(raw.rfam_hits.clone()),
        })
    }
}
<|endoftext|>


<commit_before>use super::{generator::generate_n_values, utils::did_need_in_closure_cause_panic};
use crate::{
    compiler::hir,
    database::Database,
    input::Input,
    vm::{
        tracer::Tracer,
        use_provider::DbUseProvider,
        value::{Closure, Value},
        Status, Vm,
    },
};
use std::sync::Arc;
use tokio::sync::Mutex;

pub async fn fuzz_closure(
    db: Arc<Mutex<Database>>,
    input: &Input,
    closure: Closure,
    closure_id: &hir::Id,
    mut num_instructions: usize,
) -> ClosureFuzzResult {
    loop {
        let arguments = generate_n_values(closure.num_args);
        let result = test_closure_with_args(
            db.clone(),
            closure.clone(),
            closure_id,
            arguments.clone(),
            num_instructions,
        )
        .await;

        match result {
            TestResult::DidNotFinishRunning => {
                break;
            }
            TestResult::FinishedRunningWithoutPanicking {
                num_instructions_executed,
            } => {
                num_instructions -= num_instructions_executed;
            }
            TestResult::ArgumentsDidNotFulfillNeeds {
                num_instructions_executed,
            } => {
                // This is the fuzzer's fault.
                num_instructions -= num_instructions_executed;
            }
            TestResult::InternalPanic { message, tracer } => {
                return ClosureFuzzResult::PanickedForArguments {
                    arguments,
                    message,
                    tracer,
                }
            }
        }
    }
    ClosureFuzzResult::NoProblemFound
}

pub enum ClosureFuzzResult {
    NoProblemFound,
    PanickedForArguments {
        arguments: Vec<Value>,
        message: Value,
        tracer: Tracer,
    },
}

async fn test_closure_with_args(
    db: Arc<Mutex<Database>>,
    closure: Closure,
    closure_id: &hir::Id,
    arguments: Vec<Value>,
    num_instructions: usize,
) -> TestResult {
    let mut vm = Vm::new();

    {
        let db = db.lock().await;
        println!("Starting closure {closure:?}.");
        let use_provider = DbUseProvider { db: &db };
        vm.set_up_closure_execution(&use_provider, closure, arguments);
        vm.run(&use_provider, num_instructions);
    }

    match vm.status() {
        Status::Running => TestResult::DidNotFinishRunning,
        Status::Done => TestResult::FinishedRunningWithoutPanicking {
            num_instructions_executed: vm.num_instructions_executed,
        },
        Status::Panicked(message) => {
            // If a `needs` directly inside the tested closure was not
            // satisfied, then the panic is closure's fault, but our fault.
            let db = db.lock().await;
            let is_our_fault =
                did_need_in_closure_cause_panic(&db, &closure_id, vm.tracer.log().last().unwrap());
            if is_our_fault {
                TestResult::ArgumentsDidNotFulfillNeeds {
                    num_instructions_executed: vm.num_instructions_executed,
                }
            } else {
                TestResult::InternalPanic {
                    message,
                    tracer: vm.tracer,
                }
            }
        }
    }
}
enum TestResult {
    DidNotFinishRunning,
    FinishedRunningWithoutPanicking { num_instructions_executed: usize },
    ArgumentsDidNotFulfillNeeds { num_instructions_executed: usize },
    InternalPanic { message: Value, tracer: Tracer },
}
<commit_msg>Actually respect the instruction limit in fuzz_closure<commit_after>use super::{generator::generate_n_values, utils::did_need_in_closure_cause_panic};
use crate::{
    compiler::hir,
    database::Database,
    input::Input,
    vm::{
        tracer::Tracer,
        use_provider::DbUseProvider,
        value::{Closure, Value},
        Status, Vm,
    },
};
use std::sync::Arc;
use tokio::sync::Mutex;

pub async fn fuzz_closure(
    db: Arc<Mutex<Database>>,
    input: &Input,
    closure: Closure,
    closure_id: &hir::Id,
    mut num_instructions: usize,
) -> ClosureFuzzResult {
    while num_instructions > 0 {
        let arguments = generate_n_values(closure.num_args);
        let result = test_closure_with_args(
            db.clone(),
            closure.clone(),
            closure_id,
            arguments.clone(),
            num_instructions,
        )
        .await;

        let num_instructions_executed = match result {
            TestResult::DidNotFinishRunning => {
                break;
            }
            TestResult::FinishedRunningWithoutPanicking {
                num_instructions_executed,
            } => num_instructions_executed,

            TestResult::ArgumentsDidNotFulfillNeeds {
                num_instructions_executed,
            } => num_instructions_executed,

            TestResult::InternalPanic { message, tracer } => {
                return ClosureFuzzResult::PanickedForArguments {
                    arguments,
                    message,
                    tracer,
                }
            }
        };
        num_instructions -= num_instructions_executed;
    }
    ClosureFuzzResult::NoProblemFound
}

pub enum ClosureFuzzResult {
    NoProblemFound,
    PanickedForArguments {
        arguments: Vec<Value>,
        message: Value,
        tracer: Tracer,
    },
}

async fn test_closure_with_args(
    db: Arc<Mutex<Database>>,
    closure: Closure,
    closure_id: &hir::Id,
    arguments: Vec<Value>,
    num_instructions: usize,
) -> TestResult {
    let mut vm = Vm::new();

    {
        let db = db.lock().await;
        println!("Starting closure {closure:?}.");
        let use_provider = DbUseProvider { db: &db };
        vm.set_up_closure_execution(&use_provider, closure, arguments);
        vm.run(&use_provider, num_instructions);
    }

    match vm.status() {
        Status::Running => TestResult::DidNotFinishRunning,
        Status::Done => TestResult::FinishedRunningWithoutPanicking {
            num_instructions_executed: vm.num_instructions_executed,
        },
        Status::Panicked(message) => {
            // If a `needs` directly inside the tested closure was not
            // satisfied, then the panic is closure's fault, but our fault.
            let db = db.lock().await;
            let is_our_fault =
                did_need_in_closure_cause_panic(&db, &closure_id, vm.tracer.log().last().unwrap());
            if is_our_fault {
                TestResult::ArgumentsDidNotFulfillNeeds {
                    num_instructions_executed: vm.num_instructions_executed,
                }
            } else {
                TestResult::InternalPanic {
                    message,
                    tracer: vm.tracer,
                }
            }
        }
    }
}
enum TestResult {
    DidNotFinishRunning,
    FinishedRunningWithoutPanicking { num_instructions_executed: usize },
    ArgumentsDidNotFulfillNeeds { num_instructions_executed: usize },
    InternalPanic { message: Value, tracer: Tracer },
}
<|endoftext|>


<commit_before>pub mod str {
    use std::io::{Write, Result};

    pub struct StringWriter {
        buf: Vec<u8>
    }

    impl StringWriter {
        pub fn new() -> StringWriter {
            StringWriter {
                buf: Vec::new()
            }
        }

        pub fn to_string(self) -> String {
            if let Ok(s) = String::from_utf8(self.buf) {
                s
            } else {
                String::new()
            }
        }
    }

    impl Write for StringWriter {
        fn write(&mut self, buf: &[u8]) -> Result<usize> {
            for b in buf {
                self.buf.push(*b);
            }
            Ok(buf.len())
        }

        fn flush(&mut self) -> Result<()> {
            Ok(())
        }
    }

    #[cfg(test)]
    mod test {
        use support::str::StringWriter;
        use std::io::Write;

        #[test]
        fn test_string_writer() {
            let mut sw = StringWriter::new();

            let _ = sw.write("hello".to_owned().into_bytes().as_ref());
            let _ = sw.write("world".to_owned().into_bytes().as_ref());

            let s = sw.to_string();
            assert_eq!(s, "helloworld".to_string());
        }
    }
}
<commit_msg>(feat) init buffer with a base size to reduce reallocation [#120]<commit_after>pub mod str {
    use std::io::{Write, Result};

    pub struct StringWriter {
        buf: Vec<u8>,
    }

    impl StringWriter {
        pub fn new() -> StringWriter {
            StringWriter { buf: Vec::with_capacity(8 * 1024) }
        }

        pub fn to_string(self) -> String {
            if let Ok(s) = String::from_utf8(self.buf) {
                s
            } else {
                String::new()
            }
        }
    }

    impl Write for StringWriter {
        fn write(&mut self, buf: &[u8]) -> Result<usize> {
            for b in buf {
                self.buf.push(*b);
            }
            Ok(buf.len())
        }

        fn flush(&mut self) -> Result<()> {
            Ok(())
        }
    }

    #[cfg(test)]
    mod test {
        use support::str::StringWriter;
        use std::io::Write;

        #[test]
        fn test_string_writer() {
            let mut sw = StringWriter::new();

            let _ = sw.write("hello".to_owned().into_bytes().as_ref());
            let _ = sw.write("world".to_owned().into_bytes().as_ref());

            let s = sw.to_string();
            assert_eq!(s, "helloworld".to_string());
        }
    }
}
<|endoftext|>


<commit_before>// Zinc, the bare metal stack for rust.
// Copyright 2014 Vladimir "farcaller" Pouzanov <farcaller@gmail.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::rc::Rc;
use syntax::ext::base::ExtCtxt;

use builder::{Builder, TokenString, add_node_dependency};
use node;
use super::pinmap;

pub fn attach(builder: &mut Builder, _: &mut ExtCtxt, node: Rc<node::Node>) {
  node.materializer.set(Some(verify as fn(&mut Builder, &mut ExtCtxt, Rc<node::Node>)));
  for port_node in node.subnodes().iter() {
    port_node.materializer.set(Some(verify as fn(&mut Builder, &mut ExtCtxt, Rc<node::Node>)));
    add_node_dependency(&node, port_node);
    for pin_node in port_node.subnodes().iter() {
      pin_node.materializer.set(Some(build_pin as fn(&mut Builder, &mut ExtCtxt, Rc<node::Node>)));
      add_node_dependency(port_node, pin_node);
      super::add_node_dependency_on_clock(builder, pin_node);
    }
  }
}

pub fn verify(_: &mut Builder, cx: &mut ExtCtxt, node: Rc<node::Node>) {
  node.expect_no_attributes(cx);
}

fn build_pin(builder: &mut Builder, cx: &mut ExtCtxt, node: Rc<node::Node>) {
  let port_node = node.parent.clone().unwrap().upgrade().unwrap();
  let ref port_path = port_node.path;
  let port_str = format!("Port{}", match port_path.as_str().parse::<usize>().unwrap() {
    0...4 => port_path,
    other => {
      cx.parse_sess().span_diagnostic.span_err(port_node.path_span,
          format!("unknown port `{}`, allowed values: 0...4",
              other).as_str());
      return;
    }
  });
  let port = TokenString(port_str);

  if node.name.is_none() {
    cx.parse_sess().span_diagnostic.span_err(node.name_span,
        "pin node must have a name");
    return;
  }

  let direction_str = if node.get_string_attr("function").is_some() {
    "core::option::Option::None"
  } else {
    match node.get_string_attr("direction").unwrap().as_str() {
      "out" => "core::option::Option::Some(zinc::hal::pin::Out)",
      "in"  => "core::option::Option::Some(zinc::hal::pin::In)",
      other => {
        let attr = node.get_attr("direction");
        cx.parse_sess().span_diagnostic.span_err(attr.value_span,
            format!("unknown direction `{}`, allowed values: `in`, `out`",
                other).as_str());
        return;
      }
    }
  };
  let direction = TokenString(direction_str.to_string());

  let pin_str = match node.path.as_str().parse::<usize>().unwrap() {
    0...31 => &node.path,
    other  => {
      cx.parse_sess().span_diagnostic.span_err(node.path_span,
          format!("unknown pin `{}`, allowed values: 0...31",
              other).as_str());
      return;
    }
  };

  let port_def = pinmap::port_def();
  let function_str = match node.get_string_attr("function") {
    None => "Gpio".to_string(),
    Some(fun) => {
      let pins = &port_def[port_path];
      let maybe_pin_index: usize = node.path.as_str().parse().unwrap();
      let maybe_pin: &Option<pinmap::PinDef> = pins.get(maybe_pin_index).unwrap();
      match maybe_pin {
        &None => {
          cx.parse_sess().span_diagnostic.span_err(
              node.get_attr("function").value_span,
              format!("unknown pin function `{}`, only GPIO avaliable on this pin",
                  fun).as_str());
          return;
        }
        &Some(ref pin_funcs) => {
          let maybe_func = pin_funcs.get(&fun);
          match maybe_func {
            None => {
              let avaliable: Vec<String> = pin_funcs.keys().map(|k|{k.to_string()}).collect();
              cx.parse_sess().span_diagnostic.span_err(
                  node.get_attr("function").value_span,
                  format!("unknown pin function `{}`, allowed functions: {}",
                      fun, avaliable.join(", ")).as_str());
              return;
            },
            Some(func_idx) => {
              format!("AltFunction{}", func_idx)
            }
          }
        }
      }
    }
  };

  let function = TokenString(function_str);
  let pin = TokenString(format!("{}u8", pin_str));
  let pin_name = TokenString(node.name.clone().unwrap());

  node.set_type_name("zinc::hal::lpc17xx::pin::Pin".to_string());

  let st = quote_stmt!(&*cx,
      let $pin_name = zinc::hal::lpc17xx::pin::Pin::new(
          zinc::hal::lpc17xx::pin::Port::$port,
          $pin,
          zinc::hal::lpc17xx::pin::Function::$function,
          $direction);
  ).unwrap();
  builder.add_main_statement(st);
}

#[cfg(test)]
mod test {
  use builder::Builder;
  use test_helpers::{assert_equal_source, with_parsed};

  #[test]
  fn builds_input_gpio() {
    with_parsed("
      gpio {
        0 {
          p1@1 { direction = \"in\"; }
        }
      }", |cx, failed, pt| {
      let mut builder = Builder::new(pt.clone(), cx);
      super::build_pin(&mut builder, cx, pt.get_by_name("p1").unwrap());
      assert!(unsafe{*failed} == false);
      assert!(builder.main_stmts().len() == 1);

      assert_equal_source(&builder.main_stmts()[0],
          "let p1 = zinc::hal::lpc17xx::pin::Pin::new(
               zinc::hal::lpc17xx::pin::Port::Port0,
               1u8,
               zinc::hal::lpc17xx::pin::Function::Gpio,
               core::option::Option::Some(zinc::hal::pin::In));");
    });
  }

  #[test]
  fn builds_output_gpio() {
    with_parsed("
      gpio {
        0 {
          p2@2 { direction = \"out\"; }
        }
      }", |cx, failed, pt| {
      let mut builder = Builder::new(pt.clone(), cx);
      super::build_pin(&mut builder, cx, pt.get_by_name("p2").unwrap());
      assert!(unsafe{*failed} == false);
      assert!(builder.main_stmts().len() == 1);

      assert_equal_source(&builder.main_stmts()[0],
          "let p2 = zinc::hal::lpc17xx::pin::Pin::new(
               zinc::hal::lpc17xx::pin::Port::Port0,
               2u8,
               zinc::hal::lpc17xx::pin::Function::Gpio,
               core::option::Option::Some(zinc::hal::pin::Out));");
    });
  }

  #[test]
  fn builds_altfn_gpio() {
    with_parsed("
      gpio {
        0 {
          p3@3 { direction = \"out\"; function = \"ad0_6\"; }
        }
      }", |cx, failed, pt| {
      let mut builder = Builder::new(pt.clone(), cx);
      super::build_pin(&mut builder, cx, pt.get_by_name("p3").unwrap());
      assert!(unsafe{*failed} == false);
      assert!(builder.main_stmts().len() == 1);

      assert_equal_source(&builder.main_stmts()[0],
          "let p3 = zinc::hal::lpc17xx::pin::Pin::new(
               zinc::hal::lpc17xx::pin::Port::Port0,
               3u8,
               zinc::hal::lpc17xx::pin::Function::AltFunction2,
               core::option::Option::None);");
    });
  }
}
<commit_msg>lpc17xx(pin_pt): Less cryptic unwrapping syntax<commit_after>// Zinc, the bare metal stack for rust.
// Copyright 2014 Vladimir "farcaller" Pouzanov <farcaller@gmail.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::rc::Rc;
use syntax::ext::base::ExtCtxt;

use builder::{Builder, TokenString, add_node_dependency};
use node;
use super::pinmap;

pub fn attach(builder: &mut Builder, _: &mut ExtCtxt, node: Rc<node::Node>) {
  node.materializer.set(Some(verify as fn(&mut Builder, &mut ExtCtxt, Rc<node::Node>)));
  for port_node in node.subnodes().iter() {
    port_node.materializer.set(Some(verify as fn(&mut Builder, &mut ExtCtxt, Rc<node::Node>)));
    add_node_dependency(&node, port_node);
    for pin_node in port_node.subnodes().iter() {
      pin_node.materializer.set(Some(build_pin as fn(&mut Builder, &mut ExtCtxt, Rc<node::Node>)));
      add_node_dependency(port_node, pin_node);
      super::add_node_dependency_on_clock(builder, pin_node);
    }
  }
}

pub fn verify(_: &mut Builder, cx: &mut ExtCtxt, node: Rc<node::Node>) {
  node.expect_no_attributes(cx);
}

fn build_pin(builder: &mut Builder, cx: &mut ExtCtxt, node: Rc<node::Node>) {
  let port_node = node.parent.clone().unwrap().upgrade().unwrap();
  let ref port_path = port_node.path;
  let port_str = format!("Port{}", match port_path.as_str().parse::<usize>().unwrap() {
    0...4 => port_path,
    other => {
      cx.parse_sess().span_diagnostic.span_err(port_node.path_span,
          format!("unknown port `{}`, allowed values: 0...4",
              other).as_str());
      return;
    }
  });
  let port = TokenString(port_str);

  if node.name.is_none() {
    cx.parse_sess().span_diagnostic.span_err(node.name_span,
        "pin node must have a name");
    return;
  }

  let direction_str = if node.get_string_attr("function").is_some() {
    "core::option::Option::None"
  } else {
    match node.get_string_attr("direction").unwrap().as_str() {
      "out" => "core::option::Option::Some(zinc::hal::pin::Out)",
      "in"  => "core::option::Option::Some(zinc::hal::pin::In)",
      other => {
        let attr = node.get_attr("direction");
        cx.parse_sess().span_diagnostic.span_err(attr.value_span,
            format!("unknown direction `{}`, allowed values: `in`, `out`",
                other).as_str());
        return;
      }
    }
  };
  let direction = TokenString(direction_str.to_string());

  let pin_str = match node.path.as_str().parse::<usize>().unwrap() {
    0...31 => &node.path,
    other  => {
      cx.parse_sess().span_diagnostic.span_err(node.path_span,
          format!("unknown pin `{}`, allowed values: 0...31",
              other).as_str());
      return;
    }
  };

  let port_def = pinmap::port_def();
  let function_str = match node.get_string_attr("function") {
    None => "Gpio".to_string(),
    Some(fun) => {
      let pins = &port_def[port_path];
      let pin_index: usize = node.path.as_str().parse().unwrap();
      match pins[pin_index] {
        None => {
          cx.parse_sess().span_diagnostic.span_err(
              node.get_attr("function").value_span,
              format!("unknown pin function `{}`, only GPIO avaliable on this pin",
                  fun).as_str());
          return;
        }
        Some(ref pin_funcs) => {
          let maybe_func = pin_funcs.get(&fun);
          match maybe_func {
            None => {
              let avaliable: Vec<String> = pin_funcs.keys().map(|k|{k.to_string()}).collect();
              cx.parse_sess().span_diagnostic.span_err(
                  node.get_attr("function").value_span,
                  format!("unknown pin function `{}`, allowed functions: {}",
                      fun, avaliable.join(", ")).as_str());
              return;
            },
            Some(func_idx) => {
              format!("AltFunction{}", func_idx)
            }
          }
        }
      }
    }
  };

  let function = TokenString(function_str);
  let pin = TokenString(format!("{}u8", pin_str));
  let pin_name = TokenString(node.name.clone().unwrap());

  node.set_type_name("zinc::hal::lpc17xx::pin::Pin".to_string());

  let st = quote_stmt!(&*cx,
      let $pin_name = zinc::hal::lpc17xx::pin::Pin::new(
          zinc::hal::lpc17xx::pin::Port::$port,
          $pin,
          zinc::hal::lpc17xx::pin::Function::$function,
          $direction);
  ).unwrap();
  builder.add_main_statement(st);
}

#[cfg(test)]
mod test {
  use builder::Builder;
  use test_helpers::{assert_equal_source, with_parsed};

  #[test]
  fn builds_input_gpio() {
    with_parsed("
      gpio {
        0 {
          p1@1 { direction = \"in\"; }
        }
      }", |cx, failed, pt| {
      let mut builder = Builder::new(pt.clone(), cx);
      super::build_pin(&mut builder, cx, pt.get_by_name("p1").unwrap());
      assert!(unsafe{*failed} == false);
      assert!(builder.main_stmts().len() == 1);

      assert_equal_source(&builder.main_stmts()[0],
          "let p1 = zinc::hal::lpc17xx::pin::Pin::new(
               zinc::hal::lpc17xx::pin::Port::Port0,
               1u8,
               zinc::hal::lpc17xx::pin::Function::Gpio,
               core::option::Option::Some(zinc::hal::pin::In));");
    });
  }

  #[test]
  fn builds_output_gpio() {
    with_parsed("
      gpio {
        0 {
          p2@2 { direction = \"out\"; }
        }
      }", |cx, failed, pt| {
      let mut builder = Builder::new(pt.clone(), cx);
      super::build_pin(&mut builder, cx, pt.get_by_name("p2").unwrap());
      assert!(unsafe{*failed} == false);
      assert!(builder.main_stmts().len() == 1);

      assert_equal_source(&builder.main_stmts()[0],
          "let p2 = zinc::hal::lpc17xx::pin::Pin::new(
               zinc::hal::lpc17xx::pin::Port::Port0,
               2u8,
               zinc::hal::lpc17xx::pin::Function::Gpio,
               core::option::Option::Some(zinc::hal::pin::Out));");
    });
  }

  #[test]
  fn builds_altfn_gpio() {
    with_parsed("
      gpio {
        0 {
          p3@3 { direction = \"out\"; function = \"ad0_6\"; }
        }
      }", |cx, failed, pt| {
      let mut builder = Builder::new(pt.clone(), cx);
      super::build_pin(&mut builder, cx, pt.get_by_name("p3").unwrap());
      assert!(unsafe{*failed} == false);
      assert!(builder.main_stmts().len() == 1);

      assert_equal_source(&builder.main_stmts()[0],
          "let p3 = zinc::hal::lpc17xx::pin::Pin::new(
               zinc::hal::lpc17xx::pin::Port::Port0,
               3u8,
               zinc::hal::lpc17xx::pin::Function::AltFunction2,
               core::option::Option::None);");
    });
  }
}
<|endoftext|>