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

//! A transaction can have multiple operations on state. For example, it might update values
//! for a few existing keys. Imagine that we have the following tree.
//!
//! ```text
//!                 root0
//!                 /    \
//!                /      \
//!  key1 => value11        key2 => value21
//! ```
//!
//! The next transaction updates `key1`'s value to `value12` and `key2`'s value to `value22`.
//! Let's assume we update key2 first. Then the tree becomes:
//!
//! ```text
//!                   (on disk)              (in memory)
//!                     root0                  root1'
//!                    /     \                /     \
//!                   /   ___ \ _____________/       \
//!                  /  _/     \                      \
//!                 / _/        \                      \
//!                / /           \                      \
//!   key1 => value11           key2 => value21       key2 => value22
//!      (on disk)                 (on disk)            (in memory)
//! ```
//!
//! Note that
//!   1) we created a new version of the tree with `root1'` and the new `key2` node generated;
//!   2) both `root1'` and the new `key2` node are still held in memory within a batch of nodes
//!      that will be written into db atomically.
//!
//! Next, we need to update `key1`'s value. This time we are dealing with the tree starting from
//! the new root. Part of the tree is in memory and the rest of it is in database. We'll update the
//! left child and the new root. We should
//!   1) create a new version for `key1` child.
//!   2) update `root1'` directly instead of making another version.
//! The resulting tree should look like:
//!
//! ```text
//!                   (on disk)                                     (in memory)
//!                     root0                                         root1''
//!                    /     \                                       /     \
//!                   /       \                                     /       \
//!                  /         \                                   /         \
//!                 /           \                                 /           \
//!                /             \                               /             \
//!   key1 => value11             key2 => value21  key1 => value12              key2 => value22
//!      (on disk)                   (on disk)       (in memory)                  (in memory)
//! ```
//!
//! This means that we need to be able to tell whether to create a new version of a node or to
//! update an existing node by deleting it and creating a new node directly. `TreeCache` provides
//! APIs to cache intermediate nodes and values in memory and simplify the actual tree
//! implementation.
//!
//! If we are dealing with a single-version tree, any complex tree operation can be seen as a
//! collection of the following operations:
//!   - Put a new node.
//!   - Delete a node.
//! When we apply these operations on a multi-version tree:
//!   1) Put a new node.
//!   2) When we remove a node, if the node is in the previous on-disk version, we don't need to do
//!      anything. Otherwise we delete it from the tree cache.
//! Updating node could be operated as deletion of the node followed by insertion of the updated
//! node.

#[cfg(test)]
mod tree_cache_test;

use crate::{
    metrics::DIEM_JELLYFISH_STORAGE_READS,
    node_type::{Node, NodeKey},
    NodeStats, StaleNodeIndex, TreeReader, TreeUpdateBatch,
};
use anyhow::{bail, Result};
use diem_crypto::HashValue;
use diem_types::transaction::{Version, PRE_GENESIS_VERSION};
use std::collections::{hash_map::Entry, BTreeMap, BTreeSet, HashMap, HashSet};

/// `FrozenTreeCache` is used as a field of `TreeCache` storing all the nodes and values that
/// are generated by earlier transactions so they have to be immutable. The motivation of
/// `FrozenTreeCache` is to let `TreeCache` freeze intermediate results from each transaction to
/// help commit more than one transaction in a row atomically.
struct FrozenTreeCache<V> {
    /// Immutable node_cache.
    node_cache: BTreeMap<NodeKey, Node<V>>,

    /// Immutable stale_node_index_cache.
    stale_node_index_cache: BTreeSet<StaleNodeIndex>,

    /// the stats vector including the number of new nodes, new leaves, stale nodes and stale leaves.
    node_stats: Vec<NodeStats>,

    /// Frozen root hashes after each earlier transaction.
    root_hashes: Vec<HashValue>,
}

impl<V> FrozenTreeCache<V> {
    fn new() -> Self {
        Self {
            node_cache: BTreeMap::new(),
            stale_node_index_cache: BTreeSet::new(),
            node_stats: Vec::new(),
            root_hashes: Vec::new(),
        }
    }
}

/// `TreeCache` is a in-memory cache for per-transaction updates of sparse Merkle nodes and values.
pub struct TreeCache<'a, R, V> {
    /// `NodeKey` of the current root node in cache.
    root_node_key: NodeKey,

    /// The version of the transaction to which the upcoming `put`s will be related.
    next_version: Version,

    /// Intermediate nodes keyed by node hash
    node_cache: HashMap<NodeKey, Node<V>>,

    /// # of leaves in the `node_cache`,
    num_new_leaves: usize,

    /// Partial stale log. `NodeKey` to identify the stale record.
    stale_node_index_cache: HashSet<NodeKey>,

    /// # of leaves in the `stale_node_index_cache`,
    num_stale_leaves: usize,

    /// The immutable part of this cache, which will be committed to the underlying storage.
    frozen_cache: FrozenTreeCache<V>,

    /// The underlying persistent storage.
    reader: &'a R,
}

impl<'a, R, V> TreeCache<'a, R, V>
where
    R: 'a + TreeReader<V>,
    V: crate::Value,
{
    /// Constructs a new `TreeCache` instance.
    pub fn new(reader: &'a R, next_version: Version) -> Result<Self> {
        let mut node_cache = HashMap::new();
        let root_node_key = if next_version == 0 {
            let pre_genesis_root_key = NodeKey::new_empty_path(PRE_GENESIS_VERSION);
            let pre_genesis_root = reader.get_node_option(&pre_genesis_root_key)?;

            match pre_genesis_root {
                Some(_) => {
                    // This is to support the extreme case where things really went wild,
                    // and we need to ditch the transaction history and apply a new
                    // genesis on top of an existing state db.
                    pre_genesis_root_key
                }
                None => {
                    // Hack: We need to start from an empty tree, so we insert
                    // a null node beforehand deliberately to deal with this corner case.
                    let genesis_root_key = NodeKey::new_empty_path(0);
                    node_cache.insert(genesis_root_key.clone(), Node::new_null());
                    genesis_root_key
                }
            }
        } else {
            NodeKey::new_empty_path(next_version - 1)
        };
        Ok(Self {
            node_cache,
            stale_node_index_cache: HashSet::new(),
            frozen_cache: FrozenTreeCache::new(),
            root_node_key,
            next_version,
            reader,
            num_stale_leaves: 0,
            num_new_leaves: 0,
        })
    }

    /// Gets a node with given node key. If it doesn't exist in node cache, read from `reader`.
    pub fn get_node(&self, node_key: &NodeKey) -> Result<Node<V>> {
        Ok(if let Some(node) = self.node_cache.get(node_key) {
            node.clone()
        } else if let Some(node) = self.frozen_cache.node_cache.get(node_key) {
            node.clone()
        } else {
            DIEM_JELLYFISH_STORAGE_READS.inc();
            self.reader.get_node(node_key)?
        })
    }

    /// Gets the current root node key.
    pub fn get_root_node_key(&self) -> &NodeKey {
        &self.root_node_key
    }

    /// Set roots `node_key`.
    pub fn set_root_node_key(&mut self, root_node_key: NodeKey) {
        self.root_node_key = root_node_key;
    }

    /// Puts the node with given hash as key into node_cache.
    pub fn put_node(&mut self, node_key: NodeKey, new_node: Node<V>) -> Result<()> {
        match self.node_cache.entry(node_key) {
            Entry::Vacant(o) => {
                if new_node.is_leaf() {
                    self.num_new_leaves += 1
                }
                o.insert(new_node);
            }
            Entry::Occupied(o) => bail!("Node with key {:?} already exists in NodeBatch", o.key()),
        };
        Ok(())
    }

    /// Deletes a node with given hash.
    pub fn delete_node(&mut self, old_node_key: &NodeKey, is_leaf: bool) {
        // If node cache doesn't have this node, it means the node is in the previous version of
        // the tree on the disk.
        if self.node_cache.remove(&old_node_key).is_none() {
            let is_new_entry = self.stale_node_index_cache.insert(old_node_key.clone());
            assert!(is_new_entry, "Node gets stale twice unexpectedly.");
            if is_leaf {
                self.num_stale_leaves += 1;
            }
        } else if is_leaf {
            self.num_new_leaves -= 1;
        }
    }

    /// Freezes all the contents in cache to be immutable and clear `node_cache`.
    pub fn freeze(&mut self) {
        let root_node_key = self.get_root_node_key();
        let root_hash = self
            .get_node(root_node_key)
            .unwrap_or_else(|_| unreachable!("Root node with key {:?} must exist", root_node_key))
            .hash();
        self.frozen_cache.root_hashes.push(root_hash);
        let node_stats = NodeStats {
            new_nodes: self.node_cache.len(),
            new_leaves: self.num_new_leaves,
            stale_nodes: self.stale_node_index_cache.len(),
            stale_leaves: self.num_stale_leaves,
        };
        self.frozen_cache.node_stats.push(node_stats);
        self.frozen_cache.node_cache.extend(self.node_cache.drain());
        let stale_since_version = self.next_version;
        self.frozen_cache
            .stale_node_index_cache
            .extend(
                self.stale_node_index_cache
                    .drain()
                    .map(|node_key| StaleNodeIndex {
                        stale_since_version,
                        node_key,
                    }),
            );

        // Clean up
        self.num_stale_leaves = 0;
        self.num_new_leaves = 0;

        // Prepare for the next version after freezing
        self.next_version += 1;
    }
}

impl<'a, R, V> From<TreeCache<'a, R, V>> for (Vec<HashValue>, TreeUpdateBatch<V>)
where
    R: 'a + TreeReader<V>,
{
    fn from(tree_cache: TreeCache<'a, R, V>) -> Self {
        (
            tree_cache.frozen_cache.root_hashes,
            TreeUpdateBatch {
                node_batch: tree_cache.frozen_cache.node_cache,
                stale_node_index_batch: tree_cache.frozen_cache.stale_node_index_cache,
                node_stats: tree_cache.frozen_cache.node_stats,
            },
        )
    }
}
<commit_msg>rename use crate<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! A transaction can have multiple operations on state. For example, it might update values
//! for a few existing keys. Imagine that we have the following tree.
//!
//! ```text
//!                 root0
//!                 /    \
//!                /      \
//!  key1 => value11        key2 => value21
//! ```
//!
//! The next transaction updates `key1`'s value to `value12` and `key2`'s value to `value22`.
//! Let's assume we update key2 first. Then the tree becomes:
//!
//! ```text
//!                   (on disk)              (in memory)
//!                     root0                  root1'
//!                    /     \                /     \
//!                   /   ___ \ _____________/       \
//!                  /  _/     \                      \
//!                 / _/        \                      \
//!                / /           \                      \
//!   key1 => value11           key2 => value21       key2 => value22
//!      (on disk)                 (on disk)            (in memory)
//! ```
//!
//! Note that
//!   1) we created a new version of the tree with `root1'` and the new `key2` node generated;
//!   2) both `root1'` and the new `key2` node are still held in memory within a batch of nodes
//!      that will be written into db atomically.
//!
//! Next, we need to update `key1`'s value. This time we are dealing with the tree starting from
//! the new root. Part of the tree is in memory and the rest of it is in database. We'll update the
//! left child and the new root. We should
//!   1) create a new version for `key1` child.
//!   2) update `root1'` directly instead of making another version.
//! The resulting tree should look like:
//!
//! ```text
//!                   (on disk)                                     (in memory)
//!                     root0                                         root1''
//!                    /     \                                       /     \
//!                   /       \                                     /       \
//!                  /         \                                   /         \
//!                 /           \                                 /           \
//!                /             \                               /             \
//!   key1 => value11             key2 => value21  key1 => value12              key2 => value22
//!      (on disk)                   (on disk)       (in memory)                  (in memory)
//! ```
//!
//! This means that we need to be able to tell whether to create a new version of a node or to
//! update an existing node by deleting it and creating a new node directly. `TreeCache` provides
//! APIs to cache intermediate nodes and values in memory and simplify the actual tree
//! implementation.
//!
//! If we are dealing with a single-version tree, any complex tree operation can be seen as a
//! collection of the following operations:
//!   - Put a new node.
//!   - Delete a node.
//! When we apply these operations on a multi-version tree:
//!   1) Put a new node.
//!   2) When we remove a node, if the node is in the previous on-disk version, we don't need to do
//!      anything. Otherwise we delete it from the tree cache.
//! Updating node could be operated as deletion of the node followed by insertion of the updated
//! node.

#[cfg(test)]
mod tree_cache_test;

use crate::{
    metrics::DIEM_JELLYFISH_STORAGE_READS,
    node_type::{Node, NodeKey},
    NodeBatch, NodeStats, StaleNodeIndex, StaleNodeIndexBatch, TreeReader, TreeUpdateBatch,
};
use anyhow::{bail, Result};
use diem_crypto::HashValue;
use diem_types::transaction::{Version, PRE_GENESIS_VERSION};
use std::collections::{hash_map::Entry, BTreeMap, BTreeSet, HashMap, HashSet};

/// `FrozenTreeCache` is used as a field of `TreeCache` storing all the nodes and values that
/// are generated by earlier transactions so they have to be immutable. The motivation of
/// `FrozenTreeCache` is to let `TreeCache` freeze intermediate results from each transaction to
/// help commit more than one transaction in a row atomically.
struct FrozenTreeCache<V> {
    /// Immutable node_cache.
    node_cache: NodeBatch<V>,

    /// Immutable stale_node_index_cache.
    stale_node_index_cache: StaleNodeIndexBatch,

    /// the stats vector including the number of new nodes, new leaves, stale nodes and stale leaves.
    node_stats: Vec<NodeStats>,

    /// Frozen root hashes after each earlier transaction.
    root_hashes: Vec<HashValue>,
}

impl<V> FrozenTreeCache<V> {
    fn new() -> Self {
        Self {
            node_cache: BTreeMap::new(),
            stale_node_index_cache: BTreeSet::new(),
            node_stats: Vec::new(),
            root_hashes: Vec::new(),
        }
    }
}

/// `TreeCache` is a in-memory cache for per-transaction updates of sparse Merkle nodes and values.
pub struct TreeCache<'a, R, V> {
    /// `NodeKey` of the current root node in cache.
    root_node_key: NodeKey,

    /// The version of the transaction to which the upcoming `put`s will be related.
    next_version: Version,

    /// Intermediate nodes keyed by node hash
    node_cache: HashMap<NodeKey, Node<V>>,

    /// # of leaves in the `node_cache`,
    num_new_leaves: usize,

    /// Partial stale log. `NodeKey` to identify the stale record.
    stale_node_index_cache: HashSet<NodeKey>,

    /// # of leaves in the `stale_node_index_cache`,
    num_stale_leaves: usize,

    /// The immutable part of this cache, which will be committed to the underlying storage.
    frozen_cache: FrozenTreeCache<V>,

    /// The underlying persistent storage.
    reader: &'a R,
}

impl<'a, R, V> TreeCache<'a, R, V>
where
    R: 'a + TreeReader<V>,
    V: crate::Value,
{
    /// Constructs a new `TreeCache` instance.
    pub fn new(reader: &'a R, next_version: Version) -> Result<Self> {
        let mut node_cache = HashMap::new();
        let root_node_key = if next_version == 0 {
            let pre_genesis_root_key = NodeKey::new_empty_path(PRE_GENESIS_VERSION);
            let pre_genesis_root = reader.get_node_option(&pre_genesis_root_key)?;

            match pre_genesis_root {
                Some(_) => {
                    // This is to support the extreme case where things really went wild,
                    // and we need to ditch the transaction history and apply a new
                    // genesis on top of an existing state db.
                    pre_genesis_root_key
                }
                None => {
                    // Hack: We need to start from an empty tree, so we insert
                    // a null node beforehand deliberately to deal with this corner case.
                    let genesis_root_key = NodeKey::new_empty_path(0);
                    node_cache.insert(genesis_root_key.clone(), Node::new_null());
                    genesis_root_key
                }
            }
        } else {
            NodeKey::new_empty_path(next_version - 1)
        };
        Ok(Self {
            node_cache,
            stale_node_index_cache: HashSet::new(),
            frozen_cache: FrozenTreeCache::new(),
            root_node_key,
            next_version,
            reader,
            num_stale_leaves: 0,
            num_new_leaves: 0,
        })
    }

    /// Gets a node with given node key. If it doesn't exist in node cache, read from `reader`.
    pub fn get_node(&self, node_key: &NodeKey) -> Result<Node<V>> {
        Ok(if let Some(node) = self.node_cache.get(node_key) {
            node.clone()
        } else if let Some(node) = self.frozen_cache.node_cache.get(node_key) {
            node.clone()
        } else {
            DIEM_JELLYFISH_STORAGE_READS.inc();
            self.reader.get_node(node_key)?
        })
    }

    /// Gets the current root node key.
    pub fn get_root_node_key(&self) -> &NodeKey {
        &self.root_node_key
    }

    /// Set roots `node_key`.
    pub fn set_root_node_key(&mut self, root_node_key: NodeKey) {
        self.root_node_key = root_node_key;
    }

    /// Puts the node with given hash as key into node_cache.
    pub fn put_node(&mut self, node_key: NodeKey, new_node: Node<V>) -> Result<()> {
        match self.node_cache.entry(node_key) {
            Entry::Vacant(o) => {
                if new_node.is_leaf() {
                    self.num_new_leaves += 1
                }
                o.insert(new_node);
            }
            Entry::Occupied(o) => bail!("Node with key {:?} already exists in NodeBatch", o.key()),
        };
        Ok(())
    }

    /// Deletes a node with given hash.
    pub fn delete_node(&mut self, old_node_key: &NodeKey, is_leaf: bool) {
        // If node cache doesn't have this node, it means the node is in the previous version of
        // the tree on the disk.
        if self.node_cache.remove(&old_node_key).is_none() {
            let is_new_entry = self.stale_node_index_cache.insert(old_node_key.clone());
            assert!(is_new_entry, "Node gets stale twice unexpectedly.");
            if is_leaf {
                self.num_stale_leaves += 1;
            }
        } else if is_leaf {
            self.num_new_leaves -= 1;
        }
    }

    /// Freezes all the contents in cache to be immutable and clear `node_cache`.
    pub fn freeze(&mut self) {
        let root_node_key = self.get_root_node_key();
        let root_hash = self
            .get_node(root_node_key)
            .unwrap_or_else(|_| unreachable!("Root node with key {:?} must exist", root_node_key))
            .hash();
        self.frozen_cache.root_hashes.push(root_hash);
        let node_stats = NodeStats {
            new_nodes: self.node_cache.len(),
            new_leaves: self.num_new_leaves,
            stale_nodes: self.stale_node_index_cache.len(),
            stale_leaves: self.num_stale_leaves,
        };
        self.frozen_cache.node_stats.push(node_stats);
        self.frozen_cache.node_cache.extend(self.node_cache.drain());
        let stale_since_version = self.next_version;
        self.frozen_cache
            .stale_node_index_cache
            .extend(
                self.stale_node_index_cache
                    .drain()
                    .map(|node_key| StaleNodeIndex {
                        stale_since_version,
                        node_key,
                    }),
            );

        // Clean up
        self.num_stale_leaves = 0;
        self.num_new_leaves = 0;

        // Prepare for the next version after freezing
        self.next_version += 1;
    }
}

impl<'a, R, V> From<TreeCache<'a, R, V>> for (Vec<HashValue>, TreeUpdateBatch<V>)
where
    R: 'a + TreeReader<V>,
{
    fn from(tree_cache: TreeCache<'a, R, V>) -> Self {
        (
            tree_cache.frozen_cache.root_hashes,
            TreeUpdateBatch {
                node_batch: tree_cache.frozen_cache.node_cache,
                stale_node_index_batch: tree_cache.frozen_cache.stale_node_index_cache,
                node_stats: tree_cache.frozen_cache.node_stats,
            },
        )
    }
}
<|endoftext|>


<commit_before>use serde::{
    Deserialize,
    Serialize,
};
use std::path::Path;

use serde_with::CommaSeparator;

use anyhow::Result;
use rnc_core::grouper;
use phf::phf_map;

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Precompute {
    pub id: usize,
    description: String,
    rna_type: String,
    has_coordinates: bool,
    so_rna_type: String,

    #[serde(with = "serde_with::rust::StringWithSeparator::<CommaSeparator>")]
    databases: Vec<String>,
}

#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct PrecomputeSummary {
    description: String,
    rna_type: String,
    has_coordinates: bool,
    databases: Vec<String>,
}

static SO_MAPPING: phf::Map<&'static str, &'static str> = phf_map! {
    "SO:0001263" => "SO:0000655",
    "SO:0001637" => "SO:0000252",
    "SO:0001427" => "SO:0000673",
    "SO:1001268" => "SO:0000673",
    "SO:0000122" => "SO:0000673",
    "SO:0005836" => "SO:0000673",
    "SO:0000185" => "SO:0000673",
    "SO:0000204" => "SO:0000673",
    "SO:0000140" => "SO:0000673",
    "SO:0000243" => "SO:0000673",
    "SO:0000205" => "SO:0000673",
    "SO:1001274" => "SO:0000673",
    "SO:0000726" => "SO:0000673",
    "SO:0000077" => "SO:0000673",
    "SO:0000836" => "SO:0000673",
    "SO:0000233" => "SO:0000673",
    "SO:0001267" => "SO:0000275",
    "SO:0001243" => "SO:0001244",
};

impl Precompute {
    pub fn so_rna_type(&self) -> &str {
        let given: &str = &self.so_rna_type;
        SO_MAPPING.get(given).cloned().unwrap_or(given)
    }
}

impl From<Precompute> for PrecomputeSummary {
    fn from(pre: Precompute) -> PrecomputeSummary {
        Self {
            description: pre.description,
            rna_type: pre.rna_type.replace("_", " "),
            has_coordinates: pre.has_coordinates,
            databases: pre.databases,
        }
    }
}

impl Precompute {
    pub fn has_coordiante_string(self) -> String {
        match self.has_coordinates {
            true => String::from("True"),
            false => String::from("False"),
        }
    }
}

impl grouper::HasIndex for Precompute {
    fn index(&self) -> usize {
        self.id
    }
}

pub fn group(path: &Path, max: usize, output: &Path) -> Result<()> {
    grouper::group::<Precompute>(grouper::Criteria::ExactlyOne, &path, 1, max, &output)
}
<commit_msg>Fix up another SO term<commit_after>use serde::{
    Deserialize,
    Serialize,
};
use std::path::Path;

use serde_with::CommaSeparator;

use anyhow::Result;
use rnc_core::grouper;
use phf::phf_map;

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Precompute {
    pub id: usize,
    description: String,
    rna_type: String,
    has_coordinates: bool,
    so_rna_type: String,

    #[serde(with = "serde_with::rust::StringWithSeparator::<CommaSeparator>")]
    databases: Vec<String>,
}

#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct PrecomputeSummary {
    description: String,
    rna_type: String,
    has_coordinates: bool,
    databases: Vec<String>,
}

static SO_MAPPING: phf::Map<&'static str, &'static str> = phf_map! {
    "SO:0001263" => "SO:0000655",
    "SO:0001637" => "SO:0000252",
    "SO:0001427" => "SO:0000673",
    "SO:1001268" => "SO:0000673",
    "SO:0000122" => "SO:0000673",
    "SO:0005836" => "SO:0000673",
    "SO:0000185" => "SO:0000673",
    "SO:0000204" => "SO:0000673",
    "SO:0000140" => "SO:0000673",
    "SO:0000243" => "SO:0000673",
    "SO:0000205" => "SO:0000673",
    "SO:1001274" => "SO:0000673",
    "SO:0000726" => "SO:0000673",
    "SO:0000077" => "SO:0000673",
    "SO:0000836" => "SO:0000673",
    "SO:0000233" => "SO:0000673",
    "SO:0001267" => "SO:0000275",
    "SO:0001243" => "SO:0001244",
    "SO:0001459" => "SO:0000655",
};

impl Precompute {
    pub fn so_rna_type(&self) -> &str {
        let given: &str = &self.so_rna_type;
        SO_MAPPING.get(given).cloned().unwrap_or(given)
    }
}

impl From<Precompute> for PrecomputeSummary {
    fn from(pre: Precompute) -> PrecomputeSummary {
        Self {
            description: pre.description,
            rna_type: pre.rna_type.replace("_", " "),
            has_coordinates: pre.has_coordinates,
            databases: pre.databases,
        }
    }
}

impl Precompute {
    pub fn has_coordiante_string(self) -> String {
        match self.has_coordinates {
            true => String::from("True"),
            false => String::from("False"),
        }
    }
}

impl grouper::HasIndex for Precompute {
    fn index(&self) -> usize {
        self.id
    }
}

pub fn group(path: &Path, max: usize, output: &Path) -> Result<()> {
    grouper::group::<Precompute>(grouper::Criteria::ExactlyOne, &path, 1, max, &output)
}
<|endoftext|>


<commit_before>//! User sessions.
//!
//! Actix provides a general solution for session management. Session
//! middlewares could provide different implementations which could
//! be accessed via general session api.
//!
//! By default, only cookie session backend is implemented. Other
//! backend implementations can be added.
//!
//! In general, you insert a *session* middleware and initialize it
//! , such as a `CookieSessionBackend`. To access session data,
//! [*Session*](struct.Session.html) extractor must be used. Session
//! extractor allows us to get or set session data.
//!
//! ```rust
//! use actix_web::{web, App, HttpServer, HttpResponse, Error};
//! use actix_session::{Session, CookieSession};
//!
//! fn index(session: Session) -> Result<&'static str, Error> {
//!     // access session data
//!     if let Some(count) = session.get::<i32>("counter")? {
//!         println!("SESSION value: {}", count);
//!         session.set("counter", count+1)?;
//!     } else {
//!         session.set("counter", 1)?;
//!     }
//!
//!     Ok("Welcome!")
//! }
//!
//! fn main() -> std::io::Result<()> {
//! # std::thread::spawn(||
//!     HttpServer::new(
//!         || App::new().wrap(
//!               CookieSession::signed(&[0; 32]) // <- create cookie based session middleware
//!                     .secure(false)
//!              )
//!             .service(web::resource("/").to(|| HttpResponse::Ok())))
//!         .bind("127.0.0.1:59880")?
//!         .run()
//! # );
//! # Ok(())
//! }
//! ```
use std::cell::RefCell;
use std::rc::Rc;

use actix_web::dev::{Extensions, Payload, ServiceRequest, ServiceResponse};
use actix_web::{Error, FromRequest, HttpMessage, HttpRequest};
use hashbrown::HashMap;
use serde::de::DeserializeOwned;
use serde::Serialize;
use serde_json;

#[cfg(feature = "cookie-session")]
mod cookie;
#[cfg(feature = "cookie-session")]
pub use crate::cookie::CookieSession;

/// The high-level interface you use to modify session data.
///
/// Session object could be obtained with
/// [`RequestSession::session`](trait.RequestSession.html#tymethod.session)
/// method. `RequestSession` trait is implemented for `HttpRequest`.
///
/// ```rust
/// use actix_session::Session;
/// use actix_web::*;
///
/// fn index(session: Session) -> Result<&'static str> {
///     // access session data
///     if let Some(count) = session.get::<i32>("counter")? {
///         session.set("counter", count + 1)?;
///     } else {
///         session.set("counter", 1)?;
///     }
///
///     Ok("Welcome!")
/// }
/// # fn main() {}
/// ```
pub struct Session(Rc<RefCell<SessionInner>>);

/// Helper trait that allows to get session
pub trait UserSession {
    fn get_session(&mut self) -> Session;
}

impl UserSession for HttpRequest {
    fn get_session(&mut self) -> Session {
        Session::get_session(&mut *self.extensions_mut())
    }
}

impl UserSession for ServiceRequest {
    fn get_session(&mut self) -> Session {
        Session::get_session(&mut *self.extensions_mut())
    }
}

#[derive(PartialEq, Clone, Debug)]
pub enum SessionStatus {
    Changed,
    Purged,
    Renewed,
    Unchanged
}
impl Default for SessionStatus {
    fn default() -> SessionStatus {
        SessionStatus::Unchanged
    }
}

#[derive(Default)]
struct SessionInner {
    state: HashMap<String, String>,
    pub status: SessionStatus,
}

impl Session {
    /// Get a `value` from the session.
    pub fn get<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>, Error> {
        if let Some(s) = self.0.borrow().state.get(key) {
            Ok(Some(serde_json::from_str(s)?))
        } else {
            Ok(None)
        }
    }

    /// Set a `value` from the session.
    pub fn set<T: Serialize>(&self, key: &str, value: T) -> Result<(), Error> {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Changed;
            inner
                .state
                .insert(key.to_owned(), serde_json::to_string(&value)?);
        }
        Ok(())
    }

    /// Remove value from the session.
    pub fn remove(&self, key: &str) {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Changed;
            inner.state.remove(key);
        }
    }

    /// Clear the session.
    pub fn clear(&self) {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Changed;
            inner.state.clear()
        }
    }

    /// Removes session, both client and server side.
    pub fn purge(&self) {
        let mut inner = self.0.borrow_mut();
        inner.status = SessionStatus::Purged;
        inner.state.clear();
    }

    /// Renews the session key, assigning existing session state to new key.
    pub fn renew(&self) {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Renewed;
        }
    }

    pub fn set_session(
        data: impl Iterator<Item = (String, String)>,
        req: &mut ServiceRequest,
    ) {
        let session = Session::get_session(&mut *req.extensions_mut());
        let mut inner = session.0.borrow_mut();
        inner.state.extend(data);
    }

    pub fn get_changes<B>(
        res: &mut ServiceResponse<B>,
    ) -> (SessionStatus, Option<impl Iterator<Item = (String, String)>>) {
        if let Some(s_impl) = res
            .request()
            .extensions()
            .get::<Rc<RefCell<SessionInner>>>()
        {
            let state =
                std::mem::replace(&mut s_impl.borrow_mut().state, HashMap::new());
            (s_impl.borrow().status.clone(), Some(state.into_iter()))
        } else {
            (SessionStatus::Unchanged, None)
        }
    }

    pub fn get_session(extensions: &mut Extensions) -> Session {
        if let Some(s_impl) = extensions.get::<Rc<RefCell<SessionInner>>>() {
            return Session(Rc::clone(&s_impl));
        }
        let inner = Rc::new(RefCell::new(SessionInner::default()));
        extensions.insert(inner.clone());
        Session(inner)
    }
}

/// Extractor implementation for Session type.
///
/// ```rust
/// # use actix_web::*;
/// use actix_session::Session;
///
/// fn index(session: Session) -> Result<&'static str> {
///     // access session data
///     if let Some(count) = session.get::<i32>("counter")? {
///         session.set("counter", count + 1)?;
///     } else {
///         session.set("counter", 1)?;
///     }
///
///     Ok("Welcome!")
/// }
/// # fn main() {}
/// ```
impl FromRequest for Session {
    type Error = Error;
    type Future = Result<Session, Error>;
    type Config = ();

    #[inline]
    fn from_request(req: &HttpRequest, _: &mut Payload) -> Self::Future {
        Ok(Session::get_session(&mut *req.extensions_mut()))
    }
}

#[cfg(test)]
mod tests {
    use actix_web::{test, HttpResponse};

    use super::*;

    #[test]
    fn session() {
        let mut req = test::TestRequest::default().to_srv_request();

        Session::set_session(
            vec![("key".to_string(), "\"value\"".to_string())].into_iter(),
            &mut req,
        );
        let session = Session::get_session(&mut *req.extensions_mut());
        let res = session.get::<String>("key").unwrap();
        assert_eq!(res, Some("value".to_string()));

        session.set("key2", "value2".to_string()).unwrap();
        session.remove("key");

        let mut res = req.into_response(HttpResponse::Ok().finish());
        let (_status, state) = Session::get_changes(&mut res);
        let changes: Vec<_> = state.unwrap().collect();
        assert_eq!(changes, [("key2".to_string(), "\"value2\"".to_string())]);
    }

    #[test]
    fn get_session() {
        let mut req = test::TestRequest::default().to_srv_request();

        Session::set_session(
            vec![("key".to_string(), "\"value\"".to_string())].into_iter(),
            &mut req,
        );

        let session = req.get_session();
        let res = session.get::<String>("key").unwrap();
        assert_eq!(res, Some("value".to_string()));
    }

    #[test]
    fn purge_session() {
        let mut req = test::TestRequest::default().to_srv_request();
        let session = Session::get_session(&mut *req.extensions_mut());
        assert_eq!(session.0.borrow().status, SessionStatus::Unchanged);
        session.purge();
        assert_eq!(session.0.borrow().status, SessionStatus::Purged);
    }

    #[test]
    fn renew_session() {
        let mut req = test::TestRequest::default().to_srv_request();
        let session = Session::get_session(&mut *req.extensions_mut());
        assert_eq!(session.0.borrow().status, SessionStatus::Unchanged);
        session.renew();
        assert_eq!(session.0.borrow().status, SessionStatus::Renewed);
    }
}
<commit_msg>reverting change to get_session due to side effects<commit_after>//! User sessions.
//!
//! Actix provides a general solution for session management. Session
//! middlewares could provide different implementations which could
//! be accessed via general session api.
//!
//! By default, only cookie session backend is implemented. Other
//! backend implementations can be added.
//!
//! In general, you insert a *session* middleware and initialize it
//! , such as a `CookieSessionBackend`. To access session data,
//! [*Session*](struct.Session.html) extractor must be used. Session
//! extractor allows us to get or set session data.
//!
//! ```rust
//! use actix_web::{web, App, HttpServer, HttpResponse, Error};
//! use actix_session::{Session, CookieSession};
//!
//! fn index(session: Session) -> Result<&'static str, Error> {
//!     // access session data
//!     if let Some(count) = session.get::<i32>("counter")? {
//!         println!("SESSION value: {}", count);
//!         session.set("counter", count+1)?;
//!     } else {
//!         session.set("counter", 1)?;
//!     }
//!
//!     Ok("Welcome!")
//! }
//!
//! fn main() -> std::io::Result<()> {
//! # std::thread::spawn(||
//!     HttpServer::new(
//!         || App::new().wrap(
//!               CookieSession::signed(&[0; 32]) // <- create cookie based session middleware
//!                     .secure(false)
//!              )
//!             .service(web::resource("/").to(|| HttpResponse::Ok())))
//!         .bind("127.0.0.1:59880")?
//!         .run()
//! # );
//! # Ok(())
//! }
//! ```
use std::cell::RefCell;
use std::rc::Rc;

use actix_web::dev::{Extensions, Payload, ServiceRequest, ServiceResponse};
use actix_web::{Error, FromRequest, HttpMessage, HttpRequest};
use hashbrown::HashMap;
use serde::de::DeserializeOwned;
use serde::Serialize;
use serde_json;

#[cfg(feature = "cookie-session")]
mod cookie;
#[cfg(feature = "cookie-session")]
pub use crate::cookie::CookieSession;

/// The high-level interface you use to modify session data.
///
/// Session object could be obtained with
/// [`RequestSession::session`](trait.RequestSession.html#tymethod.session)
/// method. `RequestSession` trait is implemented for `HttpRequest`.
///
/// ```rust
/// use actix_session::Session;
/// use actix_web::*;
///
/// fn index(session: Session) -> Result<&'static str> {
///     // access session data
///     if let Some(count) = session.get::<i32>("counter")? {
///         session.set("counter", count + 1)?;
///     } else {
///         session.set("counter", 1)?;
///     }
///
///     Ok("Welcome!")
/// }
/// # fn main() {}
/// ```
pub struct Session(Rc<RefCell<SessionInner>>);

/// Helper trait that allows to get session
pub trait UserSession {
    fn get_session(&mut self) -> Session;
}

impl UserSession for HttpRequest {
    fn get_session(&mut self) -> Session {
        Session::get_session(&mut *self.extensions_mut())
    }
}

impl UserSession for ServiceRequest {
    fn get_session(&mut self) -> Session {
        Session::get_session(&mut *self.extensions_mut())
    }
}

#[derive(PartialEq, Clone, Debug)]
pub enum SessionStatus {
    Changed,
    Purged,
    Renewed,
    Unchanged
}
impl Default for SessionStatus {
    fn default() -> SessionStatus {
        SessionStatus::Unchanged
    }
}

#[derive(Default)]
struct SessionInner {
    state: HashMap<String, String>,
    pub status: SessionStatus,
}

impl Session {
    /// Get a `value` from the session.
    pub fn get<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>, Error> {
        if let Some(s) = self.0.borrow().state.get(key) {
            Ok(Some(serde_json::from_str(s)?))
        } else {
            Ok(None)
        }
    }

    /// Set a `value` from the session.
    pub fn set<T: Serialize>(&self, key: &str, value: T) -> Result<(), Error> {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Changed;
            inner
                .state
                .insert(key.to_owned(), serde_json::to_string(&value)?);
        }
        Ok(())
    }

    /// Remove value from the session.
    pub fn remove(&self, key: &str) {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Changed;
            inner.state.remove(key);
        }
    }

    /// Clear the session.
    pub fn clear(&self) {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Changed;
            inner.state.clear()
        }
    }

    /// Removes session, both client and server side.
    pub fn purge(&self) {
        let mut inner = self.0.borrow_mut();
        inner.status = SessionStatus::Purged;
        inner.state.clear();
    }

    /// Renews the session key, assigning existing session state to new key.
    pub fn renew(&self) {
        let mut inner = self.0.borrow_mut();
        if inner.status != SessionStatus::Purged {
            inner.status = SessionStatus::Renewed;
        }
    }

    pub fn set_session(
        data: impl Iterator<Item = (String, String)>,
        req: &mut ServiceRequest,
    ) {
        let session = Session::get_session(&mut *req.extensions_mut());
        let mut inner = session.0.borrow_mut();
        inner.state.extend(data);
    }

    pub fn get_changes<B>(
        res: &mut ServiceResponse<B>,
    ) -> (SessionStatus, Option<impl Iterator<Item = (String, String)>>) {
        if let Some(s_impl) = res
            .request()
            .extensions()
            .get::<Rc<RefCell<SessionInner>>>()
        {
            let state =
                std::mem::replace(&mut s_impl.borrow_mut().state, HashMap::new());
            (s_impl.borrow().status.clone(), Some(state.into_iter()))
        } else {
            (SessionStatus::Unchanged, None)
        }
    }

    fn get_session(extensions: &mut Extensions) -> Session {
        if let Some(s_impl) = extensions.get::<Rc<RefCell<SessionInner>>>() {
            return Session(Rc::clone(&s_impl));
        }
        let inner = Rc::new(RefCell::new(SessionInner::default()));
        extensions.insert(inner.clone());
        Session(inner)
    }
}

/// Extractor implementation for Session type.
///
/// ```rust
/// # use actix_web::*;
/// use actix_session::Session;
///
/// fn index(session: Session) -> Result<&'static str> {
///     // access session data
///     if let Some(count) = session.get::<i32>("counter")? {
///         session.set("counter", count + 1)?;
///     } else {
///         session.set("counter", 1)?;
///     }
///
///     Ok("Welcome!")
/// }
/// # fn main() {}
/// ```
impl FromRequest for Session {
    type Error = Error;
    type Future = Result<Session, Error>;
    type Config = ();

    #[inline]
    fn from_request(req: &HttpRequest, _: &mut Payload) -> Self::Future {
        Ok(Session::get_session(&mut *req.extensions_mut()))
    }
}

#[cfg(test)]
mod tests {
    use actix_web::{test, HttpResponse};

    use super::*;

    #[test]
    fn session() {
        let mut req = test::TestRequest::default().to_srv_request();

        Session::set_session(
            vec![("key".to_string(), "\"value\"".to_string())].into_iter(),
            &mut req,
        );
        let session = Session::get_session(&mut *req.extensions_mut());
        let res = session.get::<String>("key").unwrap();
        assert_eq!(res, Some("value".to_string()));

        session.set("key2", "value2".to_string()).unwrap();
        session.remove("key");

        let mut res = req.into_response(HttpResponse::Ok().finish());
        let (_status, state) = Session::get_changes(&mut res);
        let changes: Vec<_> = state.unwrap().collect();
        assert_eq!(changes, [("key2".to_string(), "\"value2\"".to_string())]);
    }

    #[test]
    fn get_session() {
        let mut req = test::TestRequest::default().to_srv_request();

        Session::set_session(
            vec![("key".to_string(), "\"value\"".to_string())].into_iter(),
            &mut req,
        );

        let session = req.get_session();
        let res = session.get::<String>("key").unwrap();
        assert_eq!(res, Some("value".to_string()));
    }

    #[test]
    fn purge_session() {
        let mut req = test::TestRequest::default().to_srv_request();
        let session = Session::get_session(&mut *req.extensions_mut());
        assert_eq!(session.0.borrow().status, SessionStatus::Unchanged);
        session.purge();
        assert_eq!(session.0.borrow().status, SessionStatus::Purged);
    }

    #[test]
    fn renew_session() {
        let mut req = test::TestRequest::default().to_srv_request();
        let session = Session::get_session(&mut *req.extensions_mut());
        assert_eq!(session.0.borrow().status, SessionStatus::Unchanged);
        session.renew();
        assert_eq!(session.0.borrow().status, SessionStatus::Renewed);
    }
}
<|endoftext|>


<commit_before>// (C)opyleft 2013,2014 Frank Denis

/*!
 * Bloom filter for Rust
 *
 * This is a simple but fast Bloom filter implementation, that requires only
 * 2 hash functions, generated with SipHash-2-4 using randomized keys.
 */

#![desc = "A fast Bloom filter implementation."]
#![license = "BSD"]
#![crate_id="bloomfilter#0.1"]
#![crate_type = "rlib"]
#![warn(non_camel_case_types, non_uppercase_statics, unnecessary_qualification, managed_heap_memory)]

extern crate collections;
extern crate rand;

use std::cmp;
use std::hash::Hash;
use std::num;
use std::hash::sip::SipHasher;
use std::hash::Hasher;
use collections::bitv;

#[cfg(test)]
use rand::Rng;

/// Bloom filter structure
pub struct Bloom {
    bitmap: bitv::Bitv,
    bitmap_bits: u64,
    k_num: uint,
    sips: [SipHasher, ..2]
}

impl Bloom {
/// Create a new bloom filter structure.
/// bitmap_size is the size in bytes (not bits) that will be allocated in memory
/// items_count is an estimation of the maximum number of items to store.
    pub fn new(bitmap_size: uint, items_count: uint) -> Bloom {
        assert!(bitmap_size > 0u && items_count > 0u);
        let bitmap_bits = (bitmap_size as u64) * 8u64;
        let k_num = Bloom::optimal_k_num(bitmap_bits, items_count);
        let bitmap = bitv::Bitv::new(bitmap_bits as uint, false);
        let sips = [ Bloom::sip_new(), Bloom::sip_new() ];
        Bloom {
            bitmap: bitmap,
            bitmap_bits: bitmap_bits,
            k_num: k_num,
            sips: sips
        }
    }

/// Create a new bloom filter structure.
/// items_count is an estimation of the maximum number of items to store.
/// fp_p is the wanted rate of false positives, in ]0.0, 1.0[
    pub fn new_for_fp_rate(items_count: uint, fp_p: f64) -> Bloom {
        let bitmap_size = Bloom::compute_bitmap_size(items_count, fp_p);
        Bloom::new(bitmap_size, items_count)
    }

/// Compute a recommended bitmap size for items_count items
/// and a fp_p rate of false positives.
/// fp_p obviously has to be within the ]0.0, 1.0[ range.
    pub fn compute_bitmap_size(items_count: uint, fp_p: f64) -> uint {
        assert!(items_count > 0u);
        assert!(fp_p > 0.0 && fp_p < 1.0);
        let log2 = num::ln(2.0f64);
        let log2_2 = log2 * log2;
        ((items_count as f64) * num::ln(fp_p) / (-8.0 * log2_2)).ceil() as uint
    }

/// Record the presence of an item.
    pub fn set<T: Hash>(& mut self, item: T) {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            self.bitmap.set(bit_offset, true);
        }
    }

/// Check if an item is present in the set.
/// There can be false positives, but no false negatives.
    pub fn check<T: Hash>(&self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                return false;
            }
        }
        true
    }

/// Record the presence of an item in the set,
/// and return the previous state of this item.
    pub fn check_and_set<T: Hash>(&mut self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        let mut found = true;
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                found = false;
                self.bitmap.set(bit_offset, true);
            }
        }
        found
    }

    fn optimal_k_num(bitmap_bits: u64, items_count: uint) -> uint {
        let m = bitmap_bits as f64;
        let n = items_count as f64;
        let k_num = (m / n * num::ln(2.0f64).ceil()) as uint;
        cmp::max(k_num, 1)
    }

    fn bloom_hash<T: Hash>(&self, hashes: & mut [u64, ..2],
                  item: &T, k_i: uint) -> u64 {
        if k_i < 2 {
            let sip = self.sips[k_i];
            let hash = sip.hash(item);
            hashes[k_i] = hash;
            hash
        } else {
            hashes[0] + (((k_i as u64) * hashes[1]) % 0xffffffffffffffc5)
        }
    }

    fn sip_new() -> SipHasher {
        let mut rng = rand::task_rng();
        SipHasher::new_with_keys(rand::Rand::rand(& mut rng),
                                 rand::Rand::rand(& mut rng))
    }
}

#[test]
fn bloom_test_set() {
    let mut bloom = Bloom::new(10, 80);
    let key = &rand::task_rng().gen_ascii_str(16u);
    assert!(bloom.check(key) == false);
    bloom.set(&key);
    assert!(bloom.check(key.clone()) == true);
}

#[test]
fn bloom_test_check_and_set() {
    let mut bloom = Bloom::new(10, 80);
    let key = &rand::task_rng().gen_ascii_str(16u);
    assert!(bloom.check_and_set(key) == false);
    assert!(bloom.check_and_set(key.clone()) == true);
}
<commit_msg>Move away from f64::num<commit_after>// (C)opyleft 2013,2014 Frank Denis

/*!
 * Bloom filter for Rust
 *
 * This is a simple but fast Bloom filter implementation, that requires only
 * 2 hash functions, generated with SipHash-2-4 using randomized keys.
 */

#![desc = "A fast Bloom filter implementation."]
#![license = "BSD"]
#![crate_id="bloomfilter#0.1"]
#![crate_type = "rlib"]
#![warn(non_camel_case_types, non_uppercase_statics, unnecessary_qualification, managed_heap_memory)]

extern crate collections;
extern crate rand;

use std::cmp;
use std::f64;
use std::hash::Hash;
use std::hash::sip::SipHasher;
use std::hash::Hasher;
use collections::bitv;

#[cfg(test)]
use rand::Rng;

/// Bloom filter structure
pub struct Bloom {
    bitmap: bitv::Bitv,
    bitmap_bits: u64,
    k_num: uint,
    sips: [SipHasher, ..2]
}

impl Bloom {
/// Create a new bloom filter structure.
/// bitmap_size is the size in bytes (not bits) that will be allocated in memory
/// items_count is an estimation of the maximum number of items to store.
    pub fn new(bitmap_size: uint, items_count: uint) -> Bloom {
        assert!(bitmap_size > 0u && items_count > 0u);
        let bitmap_bits = (bitmap_size as u64) * 8u64;
        let k_num = Bloom::optimal_k_num(bitmap_bits, items_count);
        let bitmap = bitv::Bitv::new(bitmap_bits as uint, false);
        let sips = [ Bloom::sip_new(), Bloom::sip_new() ];
        Bloom {
            bitmap: bitmap,
            bitmap_bits: bitmap_bits,
            k_num: k_num,
            sips: sips
        }
    }

/// Create a new bloom filter structure.
/// items_count is an estimation of the maximum number of items to store.
/// fp_p is the wanted rate of false positives, in ]0.0, 1.0[
    pub fn new_for_fp_rate(items_count: uint, fp_p: f64) -> Bloom {
        let bitmap_size = Bloom::compute_bitmap_size(items_count, fp_p);
        Bloom::new(bitmap_size, items_count)
    }

/// Compute a recommended bitmap size for items_count items
/// and a fp_p rate of false positives.
/// fp_p obviously has to be within the ]0.0, 1.0[ range.
    pub fn compute_bitmap_size(items_count: uint, fp_p: f64) -> uint {
        assert!(items_count > 0u);
        assert!(fp_p > 0.0 && fp_p < 1.0);
        let log2 = f64::ln(2.0f64);
        let log2_2 = log2 * log2;
        ((items_count as f64) * f64::ln(fp_p) / (-8.0 * log2_2)).ceil() as uint
    }

/// Record the presence of an item.
    pub fn set<T: Hash>(& mut self, item: T) {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            self.bitmap.set(bit_offset, true);
        }
    }

/// Check if an item is present in the set.
/// There can be false positives, but no false negatives.
    pub fn check<T: Hash>(&self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                return false;
            }
        }
        true
    }

/// Record the presence of an item in the set,
/// and return the previous state of this item.
    pub fn check_and_set<T: Hash>(&mut self, item: T) -> bool {
        let mut hashes = [ 0u64, 0u64 ];
        let mut found = true;
        for k_i in range(0u, self.k_num) {
            let bit_offset = (self.bloom_hash(& mut hashes, &item, k_i)
                              % self.bitmap_bits) as uint;
            if self.bitmap.get(bit_offset) == false {
                found = false;
                self.bitmap.set(bit_offset, true);
            }
        }
        found
    }

    fn optimal_k_num(bitmap_bits: u64, items_count: uint) -> uint {
        let m = bitmap_bits as f64;
        let n = items_count as f64;
        let k_num = (m / n * f64::ln(2.0f64).ceil()) as uint;
        cmp::max(k_num, 1)
    }

    fn bloom_hash<T: Hash>(&self, hashes: & mut [u64, ..2],
                  item: &T, k_i: uint) -> u64 {
        if k_i < 2 {
            let sip = self.sips[k_i];
            let hash = sip.hash(item);
            hashes[k_i] = hash;
            hash
        } else {
            hashes[0] + (((k_i as u64) * hashes[1]) % 0xffffffffffffffc5)
        }
    }

    fn sip_new() -> SipHasher {
        let mut rng = rand::task_rng();
        SipHasher::new_with_keys(rand::Rand::rand(& mut rng),
                                 rand::Rand::rand(& mut rng))
    }
}

#[test]
fn bloom_test_set() {
    let mut bloom = Bloom::new(10, 80);
    let key = &rand::task_rng().gen_ascii_str(16u);
    assert!(bloom.check(key) == false);
    bloom.set(&key);
    assert!(bloom.check(key.clone()) == true);
}

#[test]
fn bloom_test_check_and_set() {
    let mut bloom = Bloom::new(10, 80);
    let key = &rand::task_rng().gen_ascii_str(16u);
    assert!(bloom.check_and_set(key) == false);
    assert!(bloom.check_and_set(key.clone()) == true);
}
<|endoftext|>


<commit_before>use std::fmt;
use world::{HEX_INNER_RADIUS, HEX_OUTER_RADIUS};
use super::{AxialType, DefaultFloat, Vector2f};

/// A 2-dimensional vector in axial coordinates. See [here][hex-blog] for more
/// information.
///
/// [hex-blog]: http://www.redblobgames.com/grids/hexagons/#coordinates
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
#[repr(C, packed)]
pub struct AxialVector {
    pub q: AxialType,
    pub r: AxialType,
}

// TODO: implement cgmath::Zero
// TODO: implement cgmath::Array
// TODO: implement cgmath::MatricSpace
// TODO: implement cgmath::VectorSpace
// TODO: implement cgmath::InnerSpace
// TODO: implement ops::{ ... }
// TODO: Add `unit_q()` and `unit_r()` (see `Vector2::unit_x()` for reference)
// For all of the above, see
// http://bjz.github.io/cgmath/cgmath/struct.Vector2.html
//

impl AxialVector {
    pub fn new(q: AxialType, r: AxialType) -> Self {
        AxialVector { q: q, r: r }
    }

    /// Returns the position of the hexagons center in the standard coordinate
    /// system using `world::{HEX_INNER_RADIUS, HEX_OUTER_RADIUS}`.
    pub fn to_real(&self) -> Vector2f {
        Vector2f {
            x: ((2 * self.q + self.r) as DefaultFloat) * HEX_INNER_RADIUS,
            y: (self.r as DefaultFloat) * (3.0 / 2.0) * HEX_OUTER_RADIUS,
        }
    }

    /// Returns the `s` component of corresponding cube coordinates. In cube
    /// coordinates 'q + r + s = 0', so saving `s` is redundant and can be
    /// calculated on the fly when needed.
    pub fn s(&self) -> AxialType {
        -self.q - self.r
    }
}

impl fmt::Debug for AxialVector {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_tuple("")
            .field(&self.q)
            .field(&self.r)
            .finish()
    }
}
<commit_msg>Added documentation to all implementations and functions in AxialVector.<commit_after>use std::fmt;
use world::{HEX_INNER_RADIUS, HEX_OUTER_RADIUS};
use super::{AxialType, DefaultFloat, Vector2f};
use math::cgmath::{VectorSpace, Zero};
use std::ops::{Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Neg, Rem,
               RemAssign, Sub, SubAssign};
use math::cgmath::prelude::{Array, MetricSpace};
use std::cmp;

/// A 2-dimensional vector in axial coordinates. See [here][hex-blog] for more
/// information.
///
/// [hex-blog]: http://www.redblobgames.com/grids/hexagons/#coordinates
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
#[repr(C, packed)]
pub struct AxialVector {
    pub q: AxialType,
    pub r: AxialType,
}

// TODO: implement cgmath::Zero
// TODO: implement cgmath::Array
// TODO: implement cgmath::MatricSpace
// TODO: implement cgmath::VectorSpace
// TODO: implement cgmath::InnerSpace
// TODO: implement ops::{ ... }
// TODO: Add `unit_q()` and `unit_r()` (see `Vector2::unit_x()` for reference)
// For all of the above, see
// http://bjz.github.io/cgmath/cgmath/struct.Vector2.html
//

/// <summary>
/// AxialVector defines a vector specificly for Axial cordinate system.
/// </summary>
impl AxialVector {
    pub fn new(q: AxialType, r: AxialType) -> Self {
        AxialVector { q: q, r: r }
    }

    /// Returns the position of the hexagons center in the standard coordinate
    /// system using `world::{HEX_INNER_RADIUS, HEX_OUTER_RADIUS}`.
    pub fn to_real(&self) -> Vector2f {
        Vector2f {
            x: ((2 * self.q + self.r) as DefaultFloat) * HEX_INNER_RADIUS,
            y: (self.r as DefaultFloat) * (3.0 / 2.0) * HEX_OUTER_RADIUS,
        }
    }

    /// Returns the `s` component of corresponding cube coordinates. In cube
    /// coordinates 'q + r + s = 0', so saving `s` is redundant and can be
    /// calculated on the fly when needed.
    pub fn s(&self) -> AxialType {
        -self.q - self.r
    }

    /// <summary>
    /// unit_q creates an default AxialVector with q:1 r:0.
    /// </summary>
    pub fn unit_q() -> AxialVector {
        AxialVector { q: 1, r: 0 }
    }

    /// <summary>
    /// unit_r creates an default AxialVector with q:0 r:1.
    /// </summary>
    pub fn unit_r() -> AxialVector {
        AxialVector { q: 0, r: 1 }
    }
}

/// <summary>
/// implements standart debug for AxialVector
/// </summary>
impl fmt::Debug for AxialVector {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_tuple("")
            .field(&self.q)
            .field(&self.r)
            .finish()
    }
}

// ********************* Basic Arithmetic (OPS) *********************

/// <summary>
/// implements Neg for AxialVector.
/// </summary>
impl Neg for AxialVector {
    type Output = AxialVector;

    /// <summary>
    /// returns a new negated AxialVector.
    /// </summary>
    fn neg(self) -> Self::Output {
        AxialVector {
            q: -self.q,
            r: -self.r,
        }
    }
}

/// <summary>
/// implements Add for AxialVector.
/// </summary>
impl Add<AxialVector> for AxialVector {
    type Output = AxialVector;

    /// <summary>
    /// adds two AxialVector and returns a new AxialVector with the result.
    /// </summary>
    /// <param name"arg2">second AxialVector to perform operation with.</param>
    fn add(self, arg2: AxialVector) -> AxialVector {
        AxialVector {
            r: self.r + arg2.r,
            q: self.q + arg2.q,
        }
    }
}

/// <summary>
/// implements AddAssign for AxialVector.
/// </summary>
impl AddAssign<AxialVector> for AxialVector {

    /// <summary>
    /// adds a Vector to itself.
    /// </summary>
    /// <param name"arg2">second AxialVector to perform operation with.</param>
    fn add_assign(&mut self, arg2: AxialVector) {
        self.r += arg2.r;
        self.q += arg2.q;

    }
}

/// <summary>
/// implements Sub for AxialVector.
/// </summary>
impl Sub<AxialVector> for AxialVector {
    type Output = AxialVector;

    /// <summary>
    /// subtracts two AxialVector and returns a new AxialVector with the result.
    /// </summary>
    /// <param name"arg2">second AxialVector to perform operation with.</param>
    fn sub(self, arg2: AxialVector) -> AxialVector {
        AxialVector {
            r: self.r - arg2.r,
            q: self.q - arg2.q,
        }
    }
}

/// <summary>
/// implements SubAssign for AxialVector.
/// </summary>
impl SubAssign<AxialVector> for AxialVector {
    /// <summary>
    /// subtracts AxialVector from itself.
    /// </summary>
    /// <param name"arg2">second AxialVector to perform operation with.</param>
    fn sub_assign(&mut self, arg2: AxialVector) {
        self.r -= arg2.r;
        self.q -= arg2.q;

    }
}

/// <summary>
/// implements Mul for AxialVector.
/// </summary>
impl Mul<AxialType> for AxialVector {
    type Output = AxialVector;

    /// <summary>
    /// multiplicates two AxialVector and returns a new AxialVector with the result.
    /// </summary>
    /// <param name"arg2">second AxialVector to perform operation with.</param>
    fn mul(self, arg2: AxialType) -> AxialVector {
        AxialVector {
            r: self.r * arg2,
            q: self.q * arg2,
        }
    }
}

/// <summary>
/// implements MulAssign for AxialVector.
/// </summary>
impl MulAssign<AxialType> for AxialVector {
    /// <summary>
    /// multiplicates AxialVector from itself.
    /// </summary>
    /// <param name"arg2">second AxialVector to perform operation with.</param>
    fn mul_assign(&mut self, arg2: AxialType) {
        self.r *= arg2;
        self.q *= arg2;

    }
}

/// <summary>
/// implements Div for AxialVector.
/// </summary>
impl Div<AxialType> for AxialVector {
    type Output = AxialVector;

    /// <summary>
    /// divides two AxialVector and returns a new AxialVector with the result.
    /// </summary>
    /// <param name"arg2">second AxialVector to perform operation with.</param>
    fn div(self, arg2: AxialType) -> AxialVector {
        AxialVector {
            r: self.r / arg2,
            q: self.q / arg2,
        }
    }
}

/// <summary>
/// implements DivAssign for AxialVector.
/// </summary>
impl DivAssign<AxialType> for AxialVector {
    /// <summary>
    /// divides AxialVector from itself.
    /// </summary>
    /// <param name"arg2">second AxialVector to perform operation with.</param>
    fn div_assign(&mut self, arg2: AxialType) {
        self.r /= arg2;
        self.q /= arg2;

    }
}

/// <summary>
/// implements Rem for AxialVector.
/// </summary>
impl Rem<AxialType> for AxialVector {
    type Output = AxialVector;

    /// <summary>
    /// calculates the modulo of AxialVector with given AxialType.
    /// </summary>
    /// <param name"arg2">second AxialType to perform operation with.</param>
    fn rem(self, arg2: AxialType) -> AxialVector {
        AxialVector {
            r: self.r % arg2,
            q: self.q % arg2,
        }
    }
}

/// <summary>
/// implements Rem for AxialVector.
/// </summary>
impl RemAssign<AxialType> for AxialVector {
    /// <summary>
    /// calculates the modulo of AxialVector with given AxialType and assignes it to itself.
    /// </summary>
    /// <param name"arg2">second AxialType to perform operation with.</param>
    fn rem_assign(&mut self, arg2: AxialType) {
        self.r %= arg2;
        self.q %= arg2;

    }
}

// ************ Metric Space ************

/// <summary>
/// implements Rem for AxialVector.
/// </summary>
impl MetricSpace for AxialVector {
    type Metric = DefaultFloat;
    /// <summary>
    /// calculates the distance betweem given AxialVector.
    /// </summary>
    /// <param name"arg2">second AxialType to perform operation with.</param>
    fn distance(self, other: AxialVector) -> DefaultFloat {
        (self.distance2(other)).sqrt()
    }
    /// <summary>
    /// calculates the distance betweem given AxialVector without using sqrt.
    /// </summary>
    /// <param name"arg2">second AxialType to perform operation with.</param>
    fn distance2(self, other: AxialVector) -> DefaultFloat {
        ((self.q - other.q).pow(2) + (self.r - other.r).pow(2)) as DefaultFloat
    }
}

// ************ Vector Space ************

/// <summary>
/// implements VectorSpace for AxialVector.
/// </summary>
impl VectorSpace for AxialVector {
    type Scalar = AxialType;
}

// ************** Zero ****************

/// <summary>
/// implements Zero for AxialVector.
/// </summary>
impl Zero for AxialVector {
    /// <summary>
    /// returns a neutral AxialVector.
    /// </summary>
    fn zero() -> AxialVector {
        AxialVector { q: 0, r: 0 }
    }

    /// <summary>
    /// checks if AxialVector is neutral element.
    /// </summary>
    fn is_zero(&self) -> bool {
        self.q == 0 && self.r == 0
    }
}

// *************** Array & Index *******************

/// <summary>
/// implements Zero for AxialVector.
/// </summary>
impl Array for AxialVector {
    type Element = AxialType;

    /// <summary>
    /// creates an AxialVector from a given value.
    /// </summary>
    /// <param name"value">value from which AxialVector is created.</param>
    fn from_value(value: Self::Element) -> Self {
        AxialVector::new(value, value)
    }

    /// <summary>
    /// sums up all elements in AxialVector and returns them.
    /// </summary>
    fn sum(self) -> Self::Element {
        self.q + self.r
    }

    /// <summary>
    /// multiplicates up all elements in AxialVector and returns them.
    /// </summary>
    fn product(self) -> Self::Element {
        self.q * self.r
    }

    /// <summary>
    /// returns smallest element in AxialVector.
    /// </summary>
    fn min(self) -> Self::Element {
        cmp::min(self.q, self.r)
    }

    /// <summary>
    /// returns largest element in AxialVector.
    /// </summary>
    fn max(self) -> Self::Element {
        cmp::max(self.q, self.r)
    }
}

/// <summary>
/// implements Index for AxialVector.
/// </summary>
impl Index<usize> for AxialVector {
    type Output = AxialType;

    /// <summary>
    /// returns AxialType of the given index number.
    /// </summary>
    /// <param name"index">element index number.</param>
    fn index(&self, index: usize) -> &Self::Output {
        let ret: &AxialType = match index {
            0 => &self.q,
            1 => &self.r,
            _ => panic!("Illegal Index Argument: was {:?}", index),
        };
        ret
    }
}

/// <summary>
/// implements IndexMut for AxialVector.
/// </summary>
impl IndexMut<usize> for AxialVector {
    /// <summary>
    /// returns a mutable AxialType of the given index number.
    /// </summary>
    /// <param name"index">element index number.</param>
    fn index_mut(&mut self, index: usize) -> &mut AxialType {
        match index {
            0 => &mut self.q,
            1 => &mut self.r,
            _ => panic!("Illegal Index Argument: was {:?}", index),
        }
    }
}

<|endoftext|>


<commit_before>// Copyright 2015-2017 Intecture Developers.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

use erased_serde::Serialize;
use errors::*;
use ExecutableProvider;
use host::*;
use pnet::datalink::interfaces;
use std::{env, process, str};
use target::{default, unix};
use telemetry::{Cpu, Os, Telemetry, TelemetryProvider, serializable};

pub struct Macos;

#[derive(Serialize, Deserialize)]
pub enum RemoteProvider {
    Available,
    Load,
}

impl <'de>ExecutableProvider<'de> for RemoteProvider {
    fn exec(&self, host: &Host) -> Result<Box<Serialize>> {
        match *self {
            RemoteProvider::Available => Ok(Box::new(Macos::available(host))),
            RemoteProvider::Load => {
                let t: serializable::Telemetry = Macos::load(host)?.into();
                Ok(Box::new(t))
            },
        }
    }
}

impl TelemetryProvider for Macos {
    fn available(host: &Host) -> bool {
        if host.is_local() {
            cfg!(target_os="macos")
        } else {
            unimplemented!();
            // let r = RemoteProvider::Available;
            // self.host.send(r).chain_err(|| ErrorKind::RemoteProvider("Telemetry", "available"))?;
            // let t: Telemetry = self.host.recv()?;
            // Ok(t)
        }
    }

    fn load(host: &Host) -> Result<Telemetry> {
        if host.is_local() {
            let cpu_vendor = try!(unix::get_sysctl_item("machdep\\.cpu\\.vendor"));
            let cpu_brand = try!(unix::get_sysctl_item("machdep\\.cpu\\.brand_string"));
            let hostname = try!(default::hostname());
            let (version_str, version_maj, version_min, version_patch) = try!(version());

            Ok(Telemetry {
                cpu: Cpu {
                    vendor: cpu_vendor,
                    brand_string: cpu_brand,
                    cores: unix::get_sysctl_item("hw\\.physicalcpu")
                                .chain_err(|| "could not resolve telemetry data")?
                                .parse::<u32>()
                                .chain_err(|| "could not resolve telemetry data")?
                },
                fs: try!(default::parse_fs(&[
                    default::FsFieldOrder::Filesystem,
                    default::FsFieldOrder::Size,
                    default::FsFieldOrder::Used,
                    default::FsFieldOrder::Available,
                    default::FsFieldOrder::Capacity,
                    default::FsFieldOrder::Blank,
                    default::FsFieldOrder::Blank,
                    default::FsFieldOrder::Blank,
                    default::FsFieldOrder::Mount,
                ])),
                hostname: hostname,
                memory: unix::get_sysctl_item("hw\\.memsize")
                             .chain_err(|| "could not resolve telemetry data")?
                             .parse::<u64>()
                             .chain_err(|| "could not resolve telemetry data")?,
                net: interfaces(),
                os: Os {
                    arch: env::consts::ARCH.into(),
                    family: "darwin".into(),
                    platform: "macos".into(),
                    version_str: version_str,
                    version_maj: version_maj,
                    version_min: version_min,
                    version_patch: version_patch
                },
            })
        } else {
            unimplemented!();
            // let r = RemoteProvider::Load;
            // self.host.send(r).chain_err(|| ErrorKind::RemoteProvider("Telemetry", "load"))?;
            // let t: Telemetry = self.host.recv()?;
            // Ok(t)
        }
    }
}

fn version() -> Result<(String, u32, u32, u32)> {
    let out = process::Command::new("sw_vers")
                               .arg("-productVersion")
                               .output()
                               .chain_err(|| ErrorKind::SystemCommand("sw_vers"))?;
    let version_str = str::from_utf8(&out.stdout)
                          .chain_err(|| ErrorKind::SystemCommandOutput("sw_vers"))?
                          .trim()
                          .to_owned();
    let (maj, min, patch) = {
        let mut parts = version_str.split('.');
        let errstr = format!("Expected OS version format `u32.u32[.u32]`, got: '{}'", version_str);
        (
            parts.next().ok_or(&*errstr)?.parse().chain_err(|| ErrorKind::SystemCommandOutput("sw_vers"))?,
            parts.next().ok_or(&*errstr)?.parse().chain_err(|| ErrorKind::SystemCommandOutput("sw_vers"))?,
            parts.next().unwrap_or("0").parse().chain_err(|| ErrorKind::SystemCommandOutput("sw_vers"))?
        )
    };
    Ok((version_str, maj, min, patch))
}
<commit_msg>Replace instances of try! with question mark operator<commit_after>// Copyright 2015-2017 Intecture Developers.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

use erased_serde::Serialize;
use errors::*;
use ExecutableProvider;
use host::*;
use pnet::datalink::interfaces;
use std::{env, process, str};
use target::{default, unix};
use telemetry::{Cpu, Os, Telemetry, TelemetryProvider, serializable};

pub struct Macos;

#[derive(Serialize, Deserialize)]
pub enum RemoteProvider {
    Available,
    Load,
}

impl <'de>ExecutableProvider<'de> for RemoteProvider {
    fn exec(&self, host: &Host) -> Result<Box<Serialize>> {
        match *self {
            RemoteProvider::Available => Ok(Box::new(Macos::available(host))),
            RemoteProvider::Load => {
                let t: serializable::Telemetry = Macos::load(host)?.into();
                Ok(Box::new(t))
            },
        }
    }
}

impl TelemetryProvider for Macos {
    fn available(host: &Host) -> bool {
        if host.is_local() {
            cfg!(target_os="macos")
        } else {
            unimplemented!();
            // let r = RemoteProvider::Available;
            // self.host.send(r).chain_err(|| ErrorKind::RemoteProvider("Telemetry", "available"))?;
            // let t: Telemetry = self.host.recv()?;
            // Ok(t)
        }
    }

    fn load(host: &Host) -> Result<Telemetry> {
        if host.is_local() {
            let cpu_vendor = unix::get_sysctl_item("machdep\\.cpu\\.vendor")?;
            let cpu_brand = unix::get_sysctl_item("machdep\\.cpu\\.brand_string")?;
            let hostname = default::hostname()?;
            let (version_str, version_maj, version_min, version_patch) = version()?;

            Ok(Telemetry {
                cpu: Cpu {
                    vendor: cpu_vendor,
                    brand_string: cpu_brand,
                    cores: unix::get_sysctl_item("hw\\.physicalcpu")
                                .chain_err(|| "could not resolve telemetry data")?
                                .parse::<u32>()
                                .chain_err(|| "could not resolve telemetry data")?
                },
                fs: default::parse_fs(&[
                    default::FsFieldOrder::Filesystem,
                    default::FsFieldOrder::Size,
                    default::FsFieldOrder::Used,
                    default::FsFieldOrder::Available,
                    default::FsFieldOrder::Capacity,
                    default::FsFieldOrder::Blank,
                    default::FsFieldOrder::Blank,
                    default::FsFieldOrder::Blank,
                    default::FsFieldOrder::Mount,
                ])?,
                hostname: hostname,
                memory: unix::get_sysctl_item("hw\\.memsize")
                             .chain_err(|| "could not resolve telemetry data")?
                             .parse::<u64>()
                             .chain_err(|| "could not resolve telemetry data")?,
                net: interfaces(),
                os: Os {
                    arch: env::consts::ARCH.into(),
                    family: "darwin".into(),
                    platform: "macos".into(),
                    version_str: version_str,
                    version_maj: version_maj,
                    version_min: version_min,
                    version_patch: version_patch
                },
            })
        } else {
            unimplemented!();
            // let r = RemoteProvider::Load;
            // self.host.send(r).chain_err(|| ErrorKind::RemoteProvider("Telemetry", "load"))?;
            // let t: Telemetry = self.host.recv()?;
            // Ok(t)
        }
    }
}

fn version() -> Result<(String, u32, u32, u32)> {
    let out = process::Command::new("sw_vers")
                               .arg("-productVersion")
                               .output()
                               .chain_err(|| ErrorKind::SystemCommand("sw_vers"))?;
    let version_str = str::from_utf8(&out.stdout)
                          .chain_err(|| ErrorKind::SystemCommandOutput("sw_vers"))?
                          .trim()
                          .to_owned();
    let (maj, min, patch) = {
        let mut parts = version_str.split('.');
        let errstr = format!("Expected OS version format `u32.u32[.u32]`, got: '{}'", version_str);
        (
            parts.next().ok_or(&*errstr)?.parse().chain_err(|| ErrorKind::SystemCommandOutput("sw_vers"))?,
            parts.next().ok_or(&*errstr)?.parse().chain_err(|| ErrorKind::SystemCommandOutput("sw_vers"))?,
            parts.next().unwrap_or("0").parse().chain_err(|| ErrorKind::SystemCommandOutput("sw_vers"))?
        )
    };
    Ok((version_str, maj, min, patch))
}
<|endoftext|>


<commit_before>//! # urdf visualization
//!
//! # Limitation
//!
//! ## Mesh
//!
//! Only `.obj` is supported by [kiss3d](https://github.com/sebcrozet/kiss3d),
//! the visualization library used this crate.
//! Other files are converted by `meshlabserver`.
//! If you add `-a` option, `assimp` is used instead of meshlab.
//! You need to install meshlab or assimp anyway.
//!
extern crate alga;
extern crate glfw;
extern crate kiss3d;
extern crate nalgebra as na;
extern crate k;
extern crate regex;
extern crate urdf_rs;
#[macro_use]
extern crate log;
extern crate rayon;
extern crate structopt;
#[macro_use]
extern crate structopt_derive;

use kiss3d::scene::SceneNode;
use kiss3d::window::Window;
use rayon::prelude::*;
use regex::Regex;
use std::collections::HashMap;
use std::fs;
use std::io;
use std::path::Path;
use std::path::PathBuf;
use std::process::Command;
use std::rc::Rc;

pub enum MeshConvert {
    Assimp,
    Meshlab,
}

pub fn get_cache_dir() -> &'static str {
    "/tmp/urdf_vis/"
}

pub fn clean_cahce_dir() -> io::Result<()> {
    fs::remove_dir_all(get_cache_dir())
}

fn create_parent_dir(new_path: &Path) -> Result<(), std::io::Error> {
    let new_parent_dir = new_path.parent().unwrap();
    if !new_parent_dir.is_dir() {
        info!("creating dir {}", new_parent_dir.to_str().unwrap());
        std::fs::create_dir_all(new_parent_dir)?;
    }
    Ok(())
}

pub fn convert_xacro_to_urdf<P>(filename: P, new_path: P) -> Result<(), std::io::Error>
where
    P: AsRef<Path>,
{
    create_parent_dir(new_path.as_ref())?;
    let output = Command::new("rosrun")
        .args(
            &[
                "xacro",
                "xacro",
                "--inorder",
                filename.as_ref().to_str().unwrap(),
                "-o",
                new_path.as_ref().to_str().unwrap(),
            ],
        )
        .output()
        .expect(
            "failed to execute xacro. install by apt-get install ros-*-xacro",
        );
    if output.status.success() {
        Ok(())
    } else {
        error!("{}", String::from_utf8(output.stderr).unwrap());
        Err(std::io::Error::new(
            std::io::ErrorKind::Other,
            "faild to xacro",
        ))
    }
}

pub fn convert_to_obj_file_by_meshlab(
    filename: &Path,
    new_path: &Path,
) -> Result<(), std::io::Error> {
    create_parent_dir(new_path)?;
    info!("converting {:?} to {:?}", filename, new_path);
    let output = Command::new("meshlabserver")
        .args(
            &[
                "-i",
                filename.to_str().unwrap(),
                "-o",
                new_path.to_str().unwrap(),
            ],
        )
        .output()
        .expect(
            "failed to execute meshlabserver. install by apt-get install meshlab",
        );
    if output.status.success() {
        Ok(())
    } else {
        Err(std::io::Error::new(
            std::io::ErrorKind::Other,
            "faild to meshlab",
        ))
    }
}

pub fn convert_to_obj_file_by_assimp(
    filename: &Path,
    new_path: &Path,
) -> Result<(), std::io::Error> {
    create_parent_dir(new_path)?;
    info!("converting {:?} to {:?}", filename, new_path);
    let output = Command::new("assimp")
        .args(
            &[
                "export",
                filename.to_str().unwrap(),
                new_path.to_str().unwrap(),
                "-ptv",
            ],
        )
        .output()
        .expect(
            "failed to execute meshlabserver. install by apt-get install assimp-utils",
        );
    if output.status.success() {
        Ok(())
    } else {
        Err(std::io::Error::new(
            std::io::ErrorKind::Other,
            "faild to assimp",
        ))
    }
}

fn rospack_find(package: &str) -> Option<String> {
    let output = Command::new("rospack")
        .arg("find")
        .arg(package)
        .output()
        .expect("rospack find failed");
    if output.status.success() {
        String::from_utf8(output.stdout)
            .map(|string| string.trim().to_string())
            .ok()
    } else {
        None
    }
}

fn expand_package_path(filename: &str) -> String {
    let re = Regex::new("^package://(\\w+)/").unwrap();
    re.replace(
        filename,
        |ma: &regex::Captures| match rospack_find(&ma[1]) {
            Some(found_path) => found_path + "/",
            None => panic!("failed to find ros package {}", &ma[1]),
        },
    )
}

fn get_cache_or_obj_path(path: &Path) -> PathBuf {
    if path.extension().unwrap() == "obj" {
        return path.to_path_buf();
    }
    let cache_path = get_cache_dir().to_string() + &path.with_extension("obj").to_str().unwrap();
    info!("cache obj path = {:?}", cache_path);
    Path::new(&cache_path).to_path_buf()
}

fn convert_mesh_if_needed(filename: &str, mesh_convert: &MeshConvert) {
    let replaced_filename = expand_package_path(filename);
    let path = Path::new(&replaced_filename);
    assert!(path.exists(), "{} not found", replaced_filename);
    let new_path = get_cache_or_obj_path(path);
    if !new_path.exists() {
        match *mesh_convert {
            MeshConvert::Assimp => convert_to_obj_file_by_assimp(path, new_path.as_path()),
            MeshConvert::Meshlab => convert_to_obj_file_by_meshlab(path, new_path.as_path()),
        }.unwrap();
    }
}

fn add_geometry(visual: &urdf_rs::Visual, window: &mut Window) -> Option<SceneNode> {
    let mut geom = match visual.geometry {
        urdf_rs::Geometry::Box { ref size } => {
            Some(window.add_cube(
                size[0] as f32,
                size[1] as f32,
                size[2] as f32,
            ))
        }
        urdf_rs::Geometry::Cylinder { radius, length } => {
            Some(window.add_cylinder(radius as f32, length as f32))
        }
        urdf_rs::Geometry::Sphere { radius } => Some(window.add_sphere(radius as f32)),
        urdf_rs::Geometry::Mesh {
            ref filename,
            scale,
        } => {
            let replaced_filename = expand_package_path(filename);
            let path = Path::new(&replaced_filename);
            assert!(path.exists(), "{} not found", replaced_filename);
            let new_path = get_cache_or_obj_path(path);
            let mtl_path = new_path.with_extension("mtl");
            // should be generated in advance
            assert!(new_path.exists());
            Some(window.add_obj(
                new_path.as_path(),
                mtl_path.as_path(),
                na::Vector3::new(
                    scale[0] as f32,
                    scale[1] as f32,
                    scale[2] as f32,
                ),
            ))
        }
    };
    let rgba = &visual.material.color.rgba;
    match geom {
        Some(ref mut obj) => obj.set_color(rgba[0] as f32, rgba[1] as f32, rgba[2] as f32),
        None => return None,
    }
    geom
}

pub struct Viewer {
    pub window: kiss3d::window::Window,
    pub urdf_robot: urdf_rs::Robot,
    pub scenes: HashMap<String, SceneNode>,
    pub arc_ball: kiss3d::camera::ArcBall,
    font_map: HashMap<i32, Rc<kiss3d::text::Font>>,
    font_data: &'static [u8],
}

impl Viewer {
    pub fn new(urdf_robot: urdf_rs::Robot) -> Viewer {
        let eye = na::Point3::new(0.5f32, 1.0, -3.0);
        let at = na::Point3::new(0.0f32, 0.25, 0.0);
        Viewer {
            window: kiss3d::window::Window::new_with_size("urdf_viewer", 1400, 1000),
            urdf_robot: urdf_robot,
            scenes: HashMap::new(),
            arc_ball: kiss3d::camera::ArcBall::new(eye, at),
            font_map: HashMap::new(),
            font_data: include_bytes!("font/Inconsolata.otf"),
        }
    }
    pub fn setup(&mut self, mesh_convert: MeshConvert) {
        self.window.set_light(kiss3d::light::Light::StickToCamera);

        let _ = self.urdf_robot
            .links
            .par_iter()
            .map(|l| if let urdf_rs::Geometry::Mesh {
                filename: ref f,
                scale: _,
            } = l.visual.geometry
            {
                convert_mesh_if_needed(f, &mesh_convert);
            })
            .count();
        for l in &self.urdf_robot.links {
            self.scenes.insert(
                l.name.to_string(),
                add_geometry(&l.visual, &mut self.window).unwrap(),
            );
        }
    }
    pub fn render(&mut self) -> bool {
        self.window.render_with_camera(&mut self.arc_ball)
    }
    pub fn update(&mut self, robot: &mut k::LinkTree<f32>) {
        for (trans, link_name) in
            robot.calc_link_transforms().iter().zip(
                robot.map_link(&|link| {
                    link.name.clone()
                }),
            )
        {
            match self.scenes.get_mut(&link_name) {
                Some(obj) => obj.set_local_transformation(*trans),
                None => {
                    println!("{} not found", link_name);
                }
            }
        }
    }
    pub fn draw_text(
        &mut self,
        text: &str,
        size: i32,
        pos: &na::Point2<f32>,
        color: &na::Point3<f32>,
    ) {
        self.window.draw_text(
            text,
            pos,
            self.font_map.entry(size).or_insert(
                kiss3d::text::Font::from_memory(
                    self.font_data,
                    size,
                ),
            ),
            color,
        );
    }
    pub fn events(&self) -> kiss3d::window::EventManager {
        self.window.events()
    }
}

pub fn convert_xacro_if_needed_and_get_path(input_path: &Path) -> Result<PathBuf, std::io::Error> {
    if input_path.extension().unwrap() == "xacro" {
        let abs_urdf_path = get_cache_dir().to_string() +
            input_path.with_extension("urdf").to_str().unwrap();
        let tmp_urdf_path = Path::new(&abs_urdf_path);
        convert_xacro_to_urdf(&input_path, &tmp_urdf_path)?;
        Ok(tmp_urdf_path.to_owned())
    } else {
        Ok(input_path.to_owned())
    }
}

#[derive(StructOpt, Debug)]
#[structopt(name = "urdf_viz", about = "Option for visualizing urdf")]
pub struct Opt {
    #[structopt(short = "a", long = "assimp",
                help = "Use assimp instead of meshlab to convert .dae to .obj for visualization")]
    pub assimp: bool,
    #[structopt(short = "c", long = "clean",
                help = "Clean the caches which is created by assimp or meshlab")]
    pub clean: bool,
    #[structopt(short = "d", long = "dof", help = "max dof for ik", default_value = "6")]
    pub ik_dof: usize,
    #[structopt(help = "Input urdf or xacro")]
    pub input_urdf_or_xacro: String,
}

impl Opt {
    pub fn get_mesh_convert_method(&self) -> MeshConvert {
        if self.assimp {
            MeshConvert::Assimp
        } else {
            MeshConvert::Meshlab
        }
    }
}
<commit_msg>Make get_cache_dir private<commit_after>//! # urdf visualization
//!
//! # Limitation
//!
//! ## Mesh
//!
//! Only `.obj` is supported by [kiss3d](https://github.com/sebcrozet/kiss3d),
//! the visualization library used this crate.
//! Other files are converted by `meshlabserver`.
//! If you add `-a` option, `assimp` is used instead of meshlab.
//! You need to install meshlab or assimp anyway.
//!
extern crate alga;
extern crate glfw;
extern crate kiss3d;
extern crate nalgebra as na;
extern crate k;
extern crate regex;
extern crate urdf_rs;
#[macro_use]
extern crate log;
extern crate rayon;
extern crate structopt;
#[macro_use]
extern crate structopt_derive;

use kiss3d::scene::SceneNode;
use kiss3d::window::Window;
use rayon::prelude::*;
use regex::Regex;
use std::collections::HashMap;
use std::fs;
use std::io;
use std::path::Path;
use std::path::PathBuf;
use std::process::Command;
use std::rc::Rc;

pub enum MeshConvert {
    Assimp,
    Meshlab,
}

fn get_cache_dir() -> &'static str {
    "/tmp/urdf_vis/"
}

pub fn clean_cahce_dir() -> io::Result<()> {
    fs::remove_dir_all(get_cache_dir())
}

fn create_parent_dir(new_path: &Path) -> Result<(), std::io::Error> {
    let new_parent_dir = new_path.parent().unwrap();
    if !new_parent_dir.is_dir() {
        info!("creating dir {}", new_parent_dir.to_str().unwrap());
        std::fs::create_dir_all(new_parent_dir)?;
    }
    Ok(())
}

pub fn convert_xacro_to_urdf<P>(filename: P, new_path: P) -> Result<(), std::io::Error>
where
    P: AsRef<Path>,
{
    create_parent_dir(new_path.as_ref())?;
    let output = Command::new("rosrun")
        .args(
            &[
                "xacro",
                "xacro",
                "--inorder",
                filename.as_ref().to_str().unwrap(),
                "-o",
                new_path.as_ref().to_str().unwrap(),
            ],
        )
        .output()
        .expect(
            "failed to execute xacro. install by apt-get install ros-*-xacro",
        );
    if output.status.success() {
        Ok(())
    } else {
        error!("{}", String::from_utf8(output.stderr).unwrap());
        Err(std::io::Error::new(
            std::io::ErrorKind::Other,
            "faild to xacro",
        ))
    }
}

pub fn convert_to_obj_file_by_meshlab(
    filename: &Path,
    new_path: &Path,
) -> Result<(), std::io::Error> {
    create_parent_dir(new_path)?;
    info!("converting {:?} to {:?}", filename, new_path);
    let output = Command::new("meshlabserver")
        .args(
            &[
                "-i",
                filename.to_str().unwrap(),
                "-o",
                new_path.to_str().unwrap(),
            ],
        )
        .output()
        .expect(
            "failed to execute meshlabserver. install by apt-get install meshlab",
        );
    if output.status.success() {
        Ok(())
    } else {
        Err(std::io::Error::new(
            std::io::ErrorKind::Other,
            "faild to meshlab",
        ))
    }
}

pub fn convert_to_obj_file_by_assimp(
    filename: &Path,
    new_path: &Path,
) -> Result<(), std::io::Error> {
    create_parent_dir(new_path)?;
    info!("converting {:?} to {:?}", filename, new_path);
    let output = Command::new("assimp")
        .args(
            &[
                "export",
                filename.to_str().unwrap(),
                new_path.to_str().unwrap(),
                "-ptv",
            ],
        )
        .output()
        .expect(
            "failed to execute meshlabserver. install by apt-get install assimp-utils",
        );
    if output.status.success() {
        Ok(())
    } else {
        Err(std::io::Error::new(
            std::io::ErrorKind::Other,
            "faild to assimp",
        ))
    }
}

fn rospack_find(package: &str) -> Option<String> {
    let output = Command::new("rospack")
        .arg("find")
        .arg(package)
        .output()
        .expect("rospack find failed");
    if output.status.success() {
        String::from_utf8(output.stdout)
            .map(|string| string.trim().to_string())
            .ok()
    } else {
        None
    }
}

fn expand_package_path(filename: &str) -> String {
    let re = Regex::new("^package://(\\w+)/").unwrap();
    re.replace(
        filename,
        |ma: &regex::Captures| match rospack_find(&ma[1]) {
            Some(found_path) => found_path + "/",
            None => panic!("failed to find ros package {}", &ma[1]),
        },
    )
}

fn get_cache_or_obj_path(path: &Path) -> PathBuf {
    if path.extension().unwrap() == "obj" {
        return path.to_path_buf();
    }
    let cache_path = get_cache_dir().to_string() + &path.with_extension("obj").to_str().unwrap();
    info!("cache obj path = {:?}", cache_path);
    Path::new(&cache_path).to_path_buf()
}

fn convert_mesh_if_needed(filename: &str, mesh_convert: &MeshConvert) {
    let replaced_filename = expand_package_path(filename);
    let path = Path::new(&replaced_filename);
    assert!(path.exists(), "{} not found", replaced_filename);
    let new_path = get_cache_or_obj_path(path);
    if !new_path.exists() {
        match *mesh_convert {
            MeshConvert::Assimp => convert_to_obj_file_by_assimp(path, new_path.as_path()),
            MeshConvert::Meshlab => convert_to_obj_file_by_meshlab(path, new_path.as_path()),
        }.unwrap();
    }
}

fn add_geometry(visual: &urdf_rs::Visual, window: &mut Window) -> Option<SceneNode> {
    let mut geom = match visual.geometry {
        urdf_rs::Geometry::Box { ref size } => {
            Some(window.add_cube(
                size[0] as f32,
                size[1] as f32,
                size[2] as f32,
            ))
        }
        urdf_rs::Geometry::Cylinder { radius, length } => {
            Some(window.add_cylinder(radius as f32, length as f32))
        }
        urdf_rs::Geometry::Sphere { radius } => Some(window.add_sphere(radius as f32)),
        urdf_rs::Geometry::Mesh {
            ref filename,
            scale,
        } => {
            let replaced_filename = expand_package_path(filename);
            let path = Path::new(&replaced_filename);
            assert!(path.exists(), "{} not found", replaced_filename);
            let new_path = get_cache_or_obj_path(path);
            let mtl_path = new_path.with_extension("mtl");
            // should be generated in advance
            assert!(new_path.exists());
            Some(window.add_obj(
                new_path.as_path(),
                mtl_path.as_path(),
                na::Vector3::new(
                    scale[0] as f32,
                    scale[1] as f32,
                    scale[2] as f32,
                ),
            ))
        }
    };
    let rgba = &visual.material.color.rgba;
    match geom {
        Some(ref mut obj) => obj.set_color(rgba[0] as f32, rgba[1] as f32, rgba[2] as f32),
        None => return None,
    }
    geom
}

pub struct Viewer {
    pub window: kiss3d::window::Window,
    pub urdf_robot: urdf_rs::Robot,
    pub scenes: HashMap<String, SceneNode>,
    pub arc_ball: kiss3d::camera::ArcBall,
    font_map: HashMap<i32, Rc<kiss3d::text::Font>>,
    font_data: &'static [u8],
}

impl Viewer {
    pub fn new(urdf_robot: urdf_rs::Robot) -> Viewer {
        let eye = na::Point3::new(0.5f32, 1.0, -3.0);
        let at = na::Point3::new(0.0f32, 0.25, 0.0);
        Viewer {
            window: kiss3d::window::Window::new_with_size("urdf_viewer", 1400, 1000),
            urdf_robot: urdf_robot,
            scenes: HashMap::new(),
            arc_ball: kiss3d::camera::ArcBall::new(eye, at),
            font_map: HashMap::new(),
            font_data: include_bytes!("font/Inconsolata.otf"),
        }
    }
    pub fn setup(&mut self, mesh_convert: MeshConvert) {
        self.window.set_light(kiss3d::light::Light::StickToCamera);

        let _ = self.urdf_robot
            .links
            .par_iter()
            .map(|l| if let urdf_rs::Geometry::Mesh {
                filename: ref f,
                scale: _,
            } = l.visual.geometry
            {
                convert_mesh_if_needed(f, &mesh_convert);
            })
            .count();
        for l in &self.urdf_robot.links {
            self.scenes.insert(
                l.name.to_string(),
                add_geometry(&l.visual, &mut self.window).unwrap(),
            );
        }
    }
    pub fn render(&mut self) -> bool {
        self.window.render_with_camera(&mut self.arc_ball)
    }
    pub fn update(&mut self, robot: &mut k::LinkTree<f32>) {
        for (trans, link_name) in
            robot.calc_link_transforms().iter().zip(
                robot.map_link(&|link| {
                    link.name.clone()
                }),
            )
        {
            match self.scenes.get_mut(&link_name) {
                Some(obj) => obj.set_local_transformation(*trans),
                None => {
                    println!("{} not found", link_name);
                }
            }
        }
    }
    pub fn draw_text(
        &mut self,
        text: &str,
        size: i32,
        pos: &na::Point2<f32>,
        color: &na::Point3<f32>,
    ) {
        self.window.draw_text(
            text,
            pos,
            self.font_map.entry(size).or_insert(
                kiss3d::text::Font::from_memory(
                    self.font_data,
                    size,
                ),
            ),
            color,
        );
    }
    pub fn events(&self) -> kiss3d::window::EventManager {
        self.window.events()
    }
}

pub fn convert_xacro_if_needed_and_get_path(input_path: &Path) -> Result<PathBuf, std::io::Error> {
    if input_path.extension().unwrap() == "xacro" {
        let abs_urdf_path = get_cache_dir().to_string() +
            input_path.with_extension("urdf").to_str().unwrap();
        let tmp_urdf_path = Path::new(&abs_urdf_path);
        convert_xacro_to_urdf(&input_path, &tmp_urdf_path)?;
        Ok(tmp_urdf_path.to_owned())
    } else {
        Ok(input_path.to_owned())
    }
}

#[derive(StructOpt, Debug)]
#[structopt(name = "urdf_viz", about = "Option for visualizing urdf")]
pub struct Opt {
    #[structopt(short = "a", long = "assimp",
                help = "Use assimp instead of meshlab to convert .dae to .obj for visualization")]
    pub assimp: bool,
    #[structopt(short = "c", long = "clean",
                help = "Clean the caches which is created by assimp or meshlab")]
    pub clean: bool,
    #[structopt(short = "d", long = "dof", help = "max dof for ik", default_value = "6")]
    pub ik_dof: usize,
    #[structopt(help = "Input urdf or xacro")]
    pub input_urdf_or_xacro: String,
}

impl Opt {
    pub fn get_mesh_convert_method(&self) -> MeshConvert {
        if self.assimp {
            MeshConvert::Assimp
        } else {
            MeshConvert::Meshlab
        }
    }
}
<|endoftext|>