<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::types::{aptos_coin_module_identifier, aptos_coin_resource_identifier};
use crate::{
    error::{ApiError, ApiResult},
    types::{
        Currency, CurrencyMetadata, MetadataRequest, NetworkIdentifier, PartialBlockIdentifier,
    },
    RosettaContext,
};
use aptos_crypto::{ValidCryptoMaterial, ValidCryptoMaterialStringExt};
use aptos_logger::debug;
use aptos_rest_client::{Account, Response};
use aptos_sdk::move_types::language_storage::{StructTag, TypeTag};
use aptos_types::{account_address::AccountAddress, chain_id::ChainId};
use futures::future::BoxFuture;
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use std::{convert::Infallible, fmt::LowerHex, future::Future, str::FromStr};
use warp::Filter;

/// The year 2000 in milliseconds, as this is the lower limit for Rosetta API implementations
pub const Y2K_MS: u64 = 946713600000;
pub const BLOCKCHAIN: &str = "aptos";

/// Checks the request network matches the server network
pub fn check_network(
    network_identifier: NetworkIdentifier,
    server_context: &RosettaContext,
) -> ApiResult<()> {
    if network_identifier.blockchain == BLOCKCHAIN
        || ChainId::from_str(network_identifier.network.trim())
            .map_err(|_| ApiError::NetworkIdentifierMismatch)?
            == server_context.chain_id
    {
        Ok(())
    } else {
        Err(ApiError::NetworkIdentifierMismatch)
    }
}

/// Attaches RosettaContext to warp paths
pub fn with_context(
    context: RosettaContext,
) -> impl Filter<Extract = (RosettaContext,), Error = Infallible> + Clone {
    warp::any().map(move || context.clone())
}

pub fn with_empty_request() -> impl Filter<Extract = (MetadataRequest,), Error = Infallible> + Clone
{
    warp::any().map(move || MetadataRequest {})
}

/// Handles a generic request to warp
pub fn handle_request<'a, F, R, Req, Resp>(
    handler: F,
) -> impl Fn(
    Req,
    RosettaContext,
) -> BoxFuture<'static, Result<warp::reply::WithStatus<warp::reply::Json>, Infallible>>
       + Clone
where
    F: FnOnce(Req, RosettaContext) -> R + Clone + Copy + Send + 'static,
    R: Future<Output = Result<Resp, ApiError>> + Send,
    Req: Deserialize<'a> + Send + 'static,
    Resp: std::fmt::Debug + Serialize,
{
    move |request, options| {
        let fut = async move {
            match handler(request, options).await {
                Ok(response) => {
                    debug!(
                        "Response: {}",
                        serde_json::to_string_pretty(&response).unwrap()
                    );
                    Ok(warp::reply::with_status(
                        warp::reply::json(&response),
                        warp::http::StatusCode::OK,
                    ))
                }
                Err(api_error) => {
                    debug!("Error: {:?}", api_error);
                    let status = api_error.status_code();
                    Ok(warp::reply::with_status(
                        warp::reply::json(&api_error.into_error()),
                        status,
                    ))
                }
            }
        };
        Box::pin(fut)
    }
}

pub async fn get_account(
    rest_client: &aptos_rest_client::Client,
    address: AccountAddress,
) -> ApiResult<Response<Account>> {
    rest_client
        .get_account(address)
        .await
        .map_err(|_| ApiError::AccountNotFound(Some(address.to_string())))
}

/// Retrieve the timestamp according ot the Rosetta spec (milliseconds)
pub fn get_timestamp(timestamp_usecs: u64) -> u64 {
    // note: timestamps are in microseconds, so we convert to milliseconds
    let mut timestamp = timestamp_usecs / 1000;

    // Rosetta doesn't like timestamps before 2000
    if timestamp < Y2K_MS {
        timestamp = Y2K_MS;
    }
    timestamp
}

/// Strips the `0x` prefix on hex strings
pub fn strip_hex_prefix(str: &str) -> &str {
    str.strip_prefix("0x").unwrap_or(str)
}

pub fn encode_bcs<T: Serialize>(obj: &T) -> ApiResult<String> {
    let bytes = bcs::to_bytes(obj)?;
    Ok(hex::encode(bytes))
}

pub fn decode_bcs<T: DeserializeOwned>(str: &str, type_name: &'static str) -> ApiResult<T> {
    let bytes = hex::decode(str)?;
    bcs::from_bytes(&bytes).map_err(|_| ApiError::deserialization_failed(type_name))
}

pub fn decode_key<T: DeserializeOwned + ValidCryptoMaterial>(
    str: &str,
    type_name: &'static str,
) -> ApiResult<T> {
    T::from_encoded_string(str).map_err(|_| ApiError::deserialization_failed(type_name))
}

const DEFAULT_COIN: &str = "APT";
const DEFAULT_DECIMALS: u8 = 8;

pub fn native_coin() -> Currency {
    Currency {
        symbol: DEFAULT_COIN.to_string(),
        decimals: DEFAULT_DECIMALS,
        metadata: Some(CurrencyMetadata {
            move_type: native_coin_tag().to_string(),
        }),
    }
}

pub fn native_coin_tag() -> TypeTag {
    TypeTag::Struct(StructTag {
        address: AccountAddress::ONE,
        module: aptos_coin_module_identifier(),
        name: aptos_coin_resource_identifier(),
        type_params: vec![],
    })
}

pub fn is_native_coin(currency: &Currency) -> ApiResult<()> {
    if currency == &native_coin() {
        Ok(())
    } else {
        Err(ApiError::UnsupportedCurrency(Some(currency.symbol.clone())))
    }
}

/// Determines which block to pull for the request
pub async fn get_block_index_from_request(
    server_context: &RosettaContext,
    partial_block_identifier: Option<PartialBlockIdentifier>,
) -> ApiResult<u64> {
    Ok(match partial_block_identifier {
        Some(PartialBlockIdentifier {
            index: Some(block_index),
            hash: Some(_),
        }) => block_index,
        // Lookup by block index
        Some(PartialBlockIdentifier {
            index: Some(block_index),
            hash: None,
        }) => block_index,
        // Lookup by block hash
        Some(PartialBlockIdentifier {
            index: None,
            hash: Some(hash),
        }) => BlockHash::from_str(&hash)?.block_height(server_context.chain_id)?,
        // Lookup latest version
        _ => {
            let response = server_context
                .rest_client()?
                .get_ledger_information()
                .await?;
            let state = response.state();

            state.block_height
        }
    })
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct BlockHash {
    chain_id: ChainId,
    block_height: u64,
}

impl BlockHash {
    pub fn new(chain_id: ChainId, block_height: u64) -> Self {
        BlockHash {
            chain_id,
            block_height,
        }
    }

    pub fn block_height(&self, expected_chain_id: ChainId) -> ApiResult<u64> {
        if expected_chain_id != self.chain_id {
            Err(ApiError::InvalidInput(Some(format!(
                "Invalid chain id in block hash {} expected {}",
                self.chain_id, expected_chain_id
            ))))
        } else {
            Ok(self.block_height)
        }
    }
}

impl FromStr for BlockHash {
    type Err = ApiError;

    fn from_str(str: &str) -> Result<Self, Self::Err> {
        let mut iter = str.split('-');

        let chain_id = if let Some(maybe_chain_id) = iter.next() {
            ChainId::from_str(maybe_chain_id).map_err(|_| {
                ApiError::InvalidInput(Some(format!(
                    "Invalid block hash, chain-id is invalid {}",
                    str
                )))
            })?
        } else {
            return Err(ApiError::InvalidInput(Some(format!(
                "Invalid block hash, missing chain-id or block height {}",
                str
            ))));
        };

        let block_height = if let Some(maybe_block_height) = iter.next() {
            u64::from_str(maybe_block_height).map_err(|_| {
                ApiError::InvalidInput(Some(format!(
                    "Invalid block hash, block height is invalid {}",
                    str
                )))
            })?
        } else {
            return Err(ApiError::InvalidInput(Some(format!(
                "Invalid block hash, missing block height {}",
                str
            ))));
        };

        if iter.next().is_some() {
            Err(ApiError::InvalidInput(Some(format!(
                "Invalid block hash, too many hyphens {}",
                str
            ))))
        } else {
            Ok(BlockHash::new(chain_id, block_height))
        }
    }
}

impl std::fmt::Display for BlockHash {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}-{}", self.chain_id, self.block_height)
    }
}

pub fn to_hex_lower<T: LowerHex>(obj: &T) -> String {
    format!("{:x}", obj)
}

#[cfg(test)]
mod test {
    use crate::common::BlockHash;
    use aptos_types::chain_id::{ChainId, NamedChain};
    use std::str::FromStr;

    #[test]
    pub fn chain_id_height_check() {
        let block_hash = BlockHash::new(ChainId::test(), 0);
        block_hash
            .block_height(ChainId::test())
            .expect("Matching chain id should work");
        block_hash
            .block_height(ChainId::new(NamedChain::MAINNET.id()))
            .expect_err("Mismatch chain id should not work");
    }

    #[test]
    pub fn chain_id_string_check() {
        let block_hash = BlockHash::new(ChainId::test(), 0);
        let parsed_block_hash =
            BlockHash::from_str(&block_hash.to_string()).expect("Should parse string");
        assert_eq!(block_hash, parsed_block_hash);
    }

    #[test]
    pub fn valid_block_hashes() {
        let valid_block_hashes: Vec<(&str, ChainId, u64)> = vec![
            ("testnet-0", ChainId::new(NamedChain::TESTNET.id()), 0),
            ("mainnet-20", ChainId::new(NamedChain::MAINNET.id()), 20),
            ("5-2", ChainId::new(5), 2),
        ];
        for (str, chain_id, height) in valid_block_hashes {
            let block_hash = BlockHash::from_str(str).expect("Valid block hash");
            assert_eq!(block_hash.block_height, height);
            assert_eq!(block_hash.chain_id, chain_id);
        }
    }

    #[test]
    pub fn invalid_block_hashes() {
        let invalid_block_hashes: Vec<&str> =
            vec!["testnet--1", "testnet", "1", "testnet-1-1", "1-mainnet"];
        for str in invalid_block_hashes {
            BlockHash::from_str(str).expect_err("Invalid block hash");
        }
    }
}
<commit_msg>[aptos-rosetta] Enforce chain id & blockchain name for APIs<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::types::{aptos_coin_module_identifier, aptos_coin_resource_identifier};
use crate::{
    error::{ApiError, ApiResult},
    types::{
        Currency, CurrencyMetadata, MetadataRequest, NetworkIdentifier, PartialBlockIdentifier,
    },
    RosettaContext,
};
use aptos_crypto::{ValidCryptoMaterial, ValidCryptoMaterialStringExt};
use aptos_logger::debug;
use aptos_rest_client::{Account, Response};
use aptos_sdk::move_types::language_storage::{StructTag, TypeTag};
use aptos_types::{account_address::AccountAddress, chain_id::ChainId};
use futures::future::BoxFuture;
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use std::{convert::Infallible, fmt::LowerHex, future::Future, str::FromStr};
use warp::Filter;

/// The year 2000 in milliseconds, as this is the lower limit for Rosetta API implementations
pub const Y2K_MS: u64 = 946713600000;
pub const BLOCKCHAIN: &str = "aptos";

/// Checks the request network matches the server network
pub fn check_network(
    network_identifier: NetworkIdentifier,
    server_context: &RosettaContext,
) -> ApiResult<()> {
    if network_identifier.blockchain == BLOCKCHAIN
        && ChainId::from_str(network_identifier.network.trim())
            .map_err(|_| ApiError::NetworkIdentifierMismatch)?
            == server_context.chain_id
    {
        Ok(())
    } else {
        Err(ApiError::NetworkIdentifierMismatch)
    }
}

/// Attaches RosettaContext to warp paths
pub fn with_context(
    context: RosettaContext,
) -> impl Filter<Extract = (RosettaContext,), Error = Infallible> + Clone {
    warp::any().map(move || context.clone())
}

pub fn with_empty_request() -> impl Filter<Extract = (MetadataRequest,), Error = Infallible> + Clone
{
    warp::any().map(move || MetadataRequest {})
}

/// Handles a generic request to warp
pub fn handle_request<'a, F, R, Req, Resp>(
    handler: F,
) -> impl Fn(
    Req,
    RosettaContext,
) -> BoxFuture<'static, Result<warp::reply::WithStatus<warp::reply::Json>, Infallible>>
       + Clone
where
    F: FnOnce(Req, RosettaContext) -> R + Clone + Copy + Send + 'static,
    R: Future<Output = Result<Resp, ApiError>> + Send,
    Req: Deserialize<'a> + Send + 'static,
    Resp: std::fmt::Debug + Serialize,
{
    move |request, options| {
        let fut = async move {
            match handler(request, options).await {
                Ok(response) => {
                    debug!(
                        "Response: {}",
                        serde_json::to_string_pretty(&response).unwrap()
                    );
                    Ok(warp::reply::with_status(
                        warp::reply::json(&response),
                        warp::http::StatusCode::OK,
                    ))
                }
                Err(api_error) => {
                    debug!("Error: {:?}", api_error);
                    let status = api_error.status_code();
                    Ok(warp::reply::with_status(
                        warp::reply::json(&api_error.into_error()),
                        status,
                    ))
                }
            }
        };
        Box::pin(fut)
    }
}

pub async fn get_account(
    rest_client: &aptos_rest_client::Client,
    address: AccountAddress,
) -> ApiResult<Response<Account>> {
    rest_client
        .get_account(address)
        .await
        .map_err(|_| ApiError::AccountNotFound(Some(address.to_string())))
}

/// Retrieve the timestamp according ot the Rosetta spec (milliseconds)
pub fn get_timestamp(timestamp_usecs: u64) -> u64 {
    // note: timestamps are in microseconds, so we convert to milliseconds
    let mut timestamp = timestamp_usecs / 1000;

    // Rosetta doesn't like timestamps before 2000
    if timestamp < Y2K_MS {
        timestamp = Y2K_MS;
    }
    timestamp
}

/// Strips the `0x` prefix on hex strings
pub fn strip_hex_prefix(str: &str) -> &str {
    str.strip_prefix("0x").unwrap_or(str)
}

pub fn encode_bcs<T: Serialize>(obj: &T) -> ApiResult<String> {
    let bytes = bcs::to_bytes(obj)?;
    Ok(hex::encode(bytes))
}

pub fn decode_bcs<T: DeserializeOwned>(str: &str, type_name: &'static str) -> ApiResult<T> {
    let bytes = hex::decode(str)?;
    bcs::from_bytes(&bytes).map_err(|_| ApiError::deserialization_failed(type_name))
}

pub fn decode_key<T: DeserializeOwned + ValidCryptoMaterial>(
    str: &str,
    type_name: &'static str,
) -> ApiResult<T> {
    T::from_encoded_string(str).map_err(|_| ApiError::deserialization_failed(type_name))
}

const DEFAULT_COIN: &str = "APT";
const DEFAULT_DECIMALS: u8 = 8;

pub fn native_coin() -> Currency {
    Currency {
        symbol: DEFAULT_COIN.to_string(),
        decimals: DEFAULT_DECIMALS,
        metadata: Some(CurrencyMetadata {
            move_type: native_coin_tag().to_string(),
        }),
    }
}

pub fn native_coin_tag() -> TypeTag {
    TypeTag::Struct(StructTag {
        address: AccountAddress::ONE,
        module: aptos_coin_module_identifier(),
        name: aptos_coin_resource_identifier(),
        type_params: vec![],
    })
}

pub fn is_native_coin(currency: &Currency) -> ApiResult<()> {
    if currency == &native_coin() {
        Ok(())
    } else {
        Err(ApiError::UnsupportedCurrency(Some(currency.symbol.clone())))
    }
}

/// Determines which block to pull for the request
pub async fn get_block_index_from_request(
    server_context: &RosettaContext,
    partial_block_identifier: Option<PartialBlockIdentifier>,
) -> ApiResult<u64> {
    Ok(match partial_block_identifier {
        Some(PartialBlockIdentifier {
            index: Some(block_index),
            hash: Some(_),
        }) => block_index,
        // Lookup by block index
        Some(PartialBlockIdentifier {
            index: Some(block_index),
            hash: None,
        }) => block_index,
        // Lookup by block hash
        Some(PartialBlockIdentifier {
            index: None,
            hash: Some(hash),
        }) => BlockHash::from_str(&hash)?.block_height(server_context.chain_id)?,
        // Lookup latest version
        _ => {
            let response = server_context
                .rest_client()?
                .get_ledger_information()
                .await?;
            let state = response.state();

            state.block_height
        }
    })
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct BlockHash {
    chain_id: ChainId,
    block_height: u64,
}

impl BlockHash {
    pub fn new(chain_id: ChainId, block_height: u64) -> Self {
        BlockHash {
            chain_id,
            block_height,
        }
    }

    pub fn block_height(&self, expected_chain_id: ChainId) -> ApiResult<u64> {
        if expected_chain_id != self.chain_id {
            Err(ApiError::InvalidInput(Some(format!(
                "Invalid chain id in block hash {} expected {}",
                self.chain_id, expected_chain_id
            ))))
        } else {
            Ok(self.block_height)
        }
    }
}

impl FromStr for BlockHash {
    type Err = ApiError;

    fn from_str(str: &str) -> Result<Self, Self::Err> {
        let mut iter = str.split('-');

        let chain_id = if let Some(maybe_chain_id) = iter.next() {
            ChainId::from_str(maybe_chain_id).map_err(|_| {
                ApiError::InvalidInput(Some(format!(
                    "Invalid block hash, chain-id is invalid {}",
                    str
                )))
            })?
        } else {
            return Err(ApiError::InvalidInput(Some(format!(
                "Invalid block hash, missing chain-id or block height {}",
                str
            ))));
        };

        let block_height = if let Some(maybe_block_height) = iter.next() {
            u64::from_str(maybe_block_height).map_err(|_| {
                ApiError::InvalidInput(Some(format!(
                    "Invalid block hash, block height is invalid {}",
                    str
                )))
            })?
        } else {
            return Err(ApiError::InvalidInput(Some(format!(
                "Invalid block hash, missing block height {}",
                str
            ))));
        };

        if iter.next().is_some() {
            Err(ApiError::InvalidInput(Some(format!(
                "Invalid block hash, too many hyphens {}",
                str
            ))))
        } else {
            Ok(BlockHash::new(chain_id, block_height))
        }
    }
}

impl std::fmt::Display for BlockHash {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}-{}", self.chain_id, self.block_height)
    }
}

pub fn to_hex_lower<T: LowerHex>(obj: &T) -> String {
    format!("{:x}", obj)
}

#[cfg(test)]
mod test {
    use crate::common::BlockHash;
    use aptos_types::chain_id::{ChainId, NamedChain};
    use std::str::FromStr;

    #[test]
    pub fn chain_id_height_check() {
        let block_hash = BlockHash::new(ChainId::test(), 0);
        block_hash
            .block_height(ChainId::test())
            .expect("Matching chain id should work");
        block_hash
            .block_height(ChainId::new(NamedChain::MAINNET.id()))
            .expect_err("Mismatch chain id should not work");
    }

    #[test]
    pub fn chain_id_string_check() {
        let block_hash = BlockHash::new(ChainId::test(), 0);
        let parsed_block_hash =
            BlockHash::from_str(&block_hash.to_string()).expect("Should parse string");
        assert_eq!(block_hash, parsed_block_hash);
    }

    #[test]
    pub fn valid_block_hashes() {
        let valid_block_hashes: Vec<(&str, ChainId, u64)> = vec![
            ("testnet-0", ChainId::new(NamedChain::TESTNET.id()), 0),
            ("mainnet-20", ChainId::new(NamedChain::MAINNET.id()), 20),
            ("5-2", ChainId::new(5), 2),
        ];
        for (str, chain_id, height) in valid_block_hashes {
            let block_hash = BlockHash::from_str(str).expect("Valid block hash");
            assert_eq!(block_hash.block_height, height);
            assert_eq!(block_hash.chain_id, chain_id);
        }
    }

    #[test]
    pub fn invalid_block_hashes() {
        let invalid_block_hashes: Vec<&str> =
            vec!["testnet--1", "testnet", "1", "testnet-1-1", "1-mainnet"];
        for str in invalid_block_hashes {
            BlockHash::from_str(str).expect_err("Invalid block hash");
        }
    }
}
<|endoftext|>


<commit_before>/// This list represents the default file types that ripgrep ships with. In
/// general, any file format is fair game, although it should generally be
/// limited to reasonably popular open formats. For other cases, you can add
/// types to each invocation of ripgrep with the '--type-add' flag.
///
/// If you would like to add or improve this list, please file a PR:
/// <https://github.com/BurntSushi/ripgrep>.
///
/// Please try to keep this list sorted lexicographically and wrapped to 79
/// columns (inclusive).
#[rustfmt::skip]
pub const DEFAULT_TYPES: &[(&str, &[&str])] = &[
    ("agda", &["*.agda", "*.lagda"]),
    ("aidl", &["*.aidl"]),
    ("amake", &["*.mk", "*.bp"]),
    ("asciidoc", &["*.adoc", "*.asc", "*.asciidoc"]),
    ("asm", &["*.asm", "*.s", "*.S"]),
    ("asp", &[
        "*.aspx", "*.aspx.cs", "*.aspx.vb", "*.ascx", "*.ascx.cs", "*.ascx.vb",
    ]),
    ("ats", &["*.ats", "*.dats", "*.sats", "*.hats"]),
    ("avro", &["*.avdl", "*.avpr", "*.avsc"]),
    ("awk", &["*.awk"]),
    ("bazel", &["*.bazel", "*.bzl", "*.BUILD", "*.bazelrc", "BUILD", "WORKSPACE"]),
    ("bitbake", &["*.bb", "*.bbappend", "*.bbclass", "*.conf", "*.inc"]),
    ("brotli", &["*.br"]),
    ("buildstream", &["*.bst"]),
    ("bzip2", &["*.bz2", "*.tbz2"]),
    ("c", &["*.[chH]", "*.[chH].in", "*.cats"]),
    ("cabal", &["*.cabal"]),
    ("cbor", &["*.cbor"]),
    ("ceylon", &["*.ceylon"]),
    ("clojure", &["*.clj", "*.cljc", "*.cljs", "*.cljx"]),
    ("cmake", &["*.cmake", "CMakeLists.txt"]),
    ("coffeescript", &["*.coffee"]),
    ("config", &["*.cfg", "*.conf", "*.config", "*.ini"]),
    ("coq", &["*.v"]),
    ("cpp", &[
        "*.[ChH]", "*.cc", "*.[ch]pp", "*.[ch]xx", "*.hh",  "*.inl",
        "*.[ChH].in", "*.cc.in", "*.[ch]pp.in", "*.[ch]xx.in", "*.hh.in",
    ]),
    ("creole", &["*.creole"]),
    ("crystal", &["Projectfile", "*.cr", "*.ecr", "shard.yml"]),
    ("cs", &["*.cs"]),
    ("csharp", &["*.cs"]),
    ("cshtml", &["*.cshtml"]),
    ("css", &["*.css", "*.scss"]),
    ("csv", &["*.csv"]),
    ("cuda", &["*.cu", "*.cuh"]),
    ("cython", &["*.pyx", "*.pxi", "*.pxd"]),
    ("d", &["*.d"]),
    ("dart", &["*.dart"]),
    ("dhall", &["*.dhall"]),
    ("diff", &["*.patch", "*.diff"]),
    ("docker", &["*Dockerfile*"]),
    ("dvc", &["Dvcfile", "*.dvc"]),
    ("ebuild", &["*.ebuild"]),
    ("edn", &["*.edn"]),
    ("elisp", &["*.el"]),
    ("elixir", &["*.ex", "*.eex", "*.exs"]),
    ("elm", &["*.elm"]),
    ("erb", &["*.erb"]),
    ("erlang", &["*.erl", "*.hrl"]),
    ("fidl", &["*.fidl"]),
    ("fish", &["*.fish"]),
    ("flatbuffers", &["*.fbs"]),
    ("fortran", &[
        "*.f", "*.F", "*.f77", "*.F77", "*.pfo",
        "*.f90", "*.F90", "*.f95", "*.F95",
    ]),
    ("fsharp", &["*.fs", "*.fsx", "*.fsi"]),
    ("fut", &[".fut"]),
    ("gap", &["*.g", "*.gap", "*.gi", "*.gd", "*.tst"]),
    ("gn", &["*.gn", "*.gni"]),
    ("go", &["*.go"]),
    ("gradle", &["*.gradle"]),
    ("groovy", &["*.groovy", "*.gradle"]),
    ("gzip", &["*.gz", "*.tgz"]),
    ("h", &["*.h", "*.hpp"]),
    ("haml", &["*.haml"]),
    ("haskell", &["*.hs", "*.lhs", "*.cpphs", "*.c2hs", "*.hsc"]),
    ("hbs", &["*.hbs"]),
    ("hs", &["*.hs", "*.lhs"]),
    ("html", &["*.htm", "*.html", "*.ejs"]),
    ("idris", &["*.idr", "*.lidr"]),
    ("janet", &["*.janet"]),
    ("java", &["*.java", "*.jsp", "*.jspx", "*.properties"]),
    ("jinja", &["*.j2", "*.jinja", "*.jinja2"]),
    ("jl", &["*.jl"]),
    ("js", &["*.js", "*.jsx", "*.vue"]),
    ("json", &["*.json", "composer.lock"]),
    ("jsonl", &["*.jsonl"]),
    ("julia", &["*.jl"]),
    ("jupyter", &["*.ipynb", "*.jpynb"]),
    ("k", &["*.k"]),
    ("kotlin", &["*.kt", "*.kts"]),
    ("less", &["*.less"]),
    ("license", &[
        // General
        "COPYING", "COPYING[.-]*",
        "COPYRIGHT", "COPYRIGHT[.-]*",
        "EULA", "EULA[.-]*",
        "licen[cs]e", "licen[cs]e.*",
        "LICEN[CS]E", "LICEN[CS]E[.-]*", "*[.-]LICEN[CS]E*",
        "NOTICE", "NOTICE[.-]*",
        "PATENTS", "PATENTS[.-]*",
        "UNLICEN[CS]E", "UNLICEN[CS]E[.-]*",
        // GPL (gpl.txt, etc.)
        "agpl[.-]*",
        "gpl[.-]*",
        "lgpl[.-]*",
        // Other license-specific (APACHE-2.0.txt, etc.)
        "AGPL-*[0-9]*",
        "APACHE-*[0-9]*",
        "BSD-*[0-9]*",
        "CC-BY-*",
        "GFDL-*[0-9]*",
        "GNU-*[0-9]*",
        "GPL-*[0-9]*",
        "LGPL-*[0-9]*",
        "MIT-*[0-9]*",
        "MPL-*[0-9]*",
        "OFL-*[0-9]*",
    ]),
    ("lisp", &["*.el", "*.jl", "*.lisp", "*.lsp", "*.sc", "*.scm"]),
    ("lock", &["*.lock", "package-lock.json"]),
    ("log", &["*.log"]),
    ("lua", &["*.lua"]),
    ("lz4", &["*.lz4"]),
    ("lzma", &["*.lzma"]),
    ("m4", &["*.ac", "*.m4"]),
    ("make", &[
        "[Gg][Nn][Uu]makefile", "[Mm]akefile",
        "[Gg][Nn][Uu]makefile.am", "[Mm]akefile.am",
        "[Gg][Nn][Uu]makefile.in", "[Mm]akefile.in",
        "*.mk", "*.mak"
    ]),
    ("mako", &["*.mako", "*.mao"]),
    ("man", &["*.[0-9lnpx]", "*.[0-9][cEFMmpSx]"]),
    ("markdown", &["*.markdown", "*.md", "*.mdown", "*.mkdn"]),
    ("matlab", &["*.m"]),
    ("md", &["*.markdown", "*.md", "*.mdown", "*.mkdn"]),
    ("meson", &["meson.build", "meson_options.txt"]),
    ("minified", &["*.min.html", "*.min.css", "*.min.js"]),
    ("mint", &["*.mint"]),
    ("mk", &["mkfile"]),
    ("ml", &["*.ml"]),
    ("msbuild", &[
        "*.csproj", "*.fsproj", "*.vcxproj", "*.proj", "*.props", "*.targets",
    ]),
    ("nim", &["*.nim", "*.nimf", "*.nimble", "*.nims"]),
    ("nix", &["*.nix"]),
    ("objc", &["*.h", "*.m"]),
    ("objcpp", &["*.h", "*.mm"]),
    ("ocaml", &["*.ml", "*.mli", "*.mll", "*.mly"]),
    ("org", &["*.org", "*.org_archive"]),
    ("pascal", &["*.pas", "*.dpr", "*.lpr", "*.pp", "*.inc"]),
    ("pdf", &["*.pdf"]),
    ("perl", &["*.perl", "*.pl", "*.PL", "*.plh", "*.plx", "*.pm", "*.t"]),
    ("php", &["*.php", "*.php3", "*.php4", "*.php5", "*.phtml"]),
    ("po", &["*.po"]),
    ("pod", &["*.pod"]),
    ("postscript", &["*.eps", "*.ps"]),
    ("protobuf", &["*.proto"]),
    ("ps", &["*.cdxml", "*.ps1", "*.ps1xml", "*.psd1", "*.psm1"]),
    ("puppet", &["*.erb", "*.pp", "*.rb"]),
    ("purs", &["*.purs"]),
    ("py", &["*.py"]),
    ("qmake", &["*.pro", "*.pri", "*.prf"]),
    ("qml", &["*.qml"]),
    ("r", &["*.R", "*.r", "*.Rmd", "*.Rnw"]),
    ("racket", &["*.rkt"]),
    ("rdoc", &["*.rdoc"]),
    ("readme", &["README*", "*README"]),
    ("red", &["*.r", "*.red", "*.reds"]),
    ("robot", &["*.robot"]),
    ("rst", &["*.rst"]),
    ("ruby", &[
        // Idiomatic files
        "config.ru", "Gemfile", ".irbrc", "Rakefile",
        // Extensions
        "*.gemspec", "*.rb", "*.rbw"
    ]),
    ("rust", &["*.rs"]),
    ("sass", &["*.sass", "*.scss"]),
    ("scala", &["*.scala", "*.sbt"]),
    ("sh", &[
        // Portable/misc. init files
        ".login", ".logout", ".profile", "profile",
        // bash-specific init files
        ".bash_login", "bash_login",
        ".bash_logout", "bash_logout",
        ".bash_profile", "bash_profile",
        ".bashrc", "bashrc", "*.bashrc",
        // csh-specific init files
        ".cshrc", "*.cshrc",
        // ksh-specific init files
        ".kshrc", "*.kshrc",
        // tcsh-specific init files
        ".tcshrc",
        // zsh-specific init files
        ".zshenv", "zshenv",
        ".zlogin", "zlogin",
        ".zlogout", "zlogout",
        ".zprofile", "zprofile",
        ".zshrc", "zshrc",
        // Extensions
        "*.bash", "*.csh", "*.ksh", "*.sh", "*.tcsh", "*.zsh",
    ]),
    ("slim", &["*.skim", "*.slim", "*.slime"]),
    ("smarty", &["*.tpl"]),
    ("sml", &["*.sml", "*.sig"]),
    ("soy", &["*.soy"]),
    ("spark", &["*.spark"]),
    ("spec", &["*.spec"]),
    ("sql", &["*.sql", "*.psql"]),
    ("stylus", &["*.styl"]),
    ("sv", &["*.v", "*.vg", "*.sv", "*.svh", "*.h"]),
    ("svg", &["*.svg"]),
    ("swift", &["*.swift"]),
    ("swig", &["*.def", "*.i"]),
    ("systemd", &[
        "*.automount", "*.conf", "*.device", "*.link", "*.mount", "*.path",
        "*.scope", "*.service", "*.slice", "*.socket", "*.swap", "*.target",
        "*.timer",
    ]),
    ("taskpaper", &["*.taskpaper"]),
    ("tcl", &["*.tcl"]),
    ("tex", &["*.tex", "*.ltx", "*.cls", "*.sty", "*.bib", "*.dtx", "*.ins"]),
    ("texinfo", &["*.texi"]),
    ("textile", &["*.textile"]),
    ("tf", &["*.tf"]),
    ("thrift", &["*.thrift"]),
    ("toml", &["*.toml", "Cargo.lock"]),
    ("ts", &["*.ts", "*.tsx"]),
    ("twig", &["*.twig"]),
    ("txt", &["*.txt"]),
    ("typoscript", &["*.typoscript", "*.ts"]),
    ("vala", &["*.vala"]),
    ("vb", &["*.vb"]),
    ("vcl", &["*.vcl"]),
    ("verilog", &["*.v", "*.vh", "*.sv", "*.svh"]),
    ("vhdl", &["*.vhd", "*.vhdl"]),
    ("vim", &["*.vim"]),
    ("vimscript", &["*.vim"]),
    ("webidl", &["*.idl", "*.webidl", "*.widl"]),
    ("wiki", &["*.mediawiki", "*.wiki"]),
    ("xml", &[
        "*.xml", "*.xml.dist", "*.dtd", "*.xsl", "*.xslt", "*.xsd", "*.xjb",
        "*.rng", "*.sch", "*.xhtml",
    ]),
    ("xz", &["*.xz", "*.txz"]),
    ("yacc", &["*.y"]),
    ("yaml", &["*.yaml", "*.yml"]),
    ("yang", &["*.yang"]),
    ("z", &["*.Z"]),
    ("zig", &["*.zig"]),
    ("zsh", &[
        ".zshenv", "zshenv",
        ".zlogin", "zlogin",
        ".zlogout", "zlogout",
        ".zprofile", "zprofile",
        ".zshrc", "zshrc",
        "*.zsh",
    ]),
    ("zstd", &["*.zst", "*.zstd"]),
];
<commit_msg>ignore/types: fix futhark type extension<commit_after>/// This list represents the default file types that ripgrep ships with. In
/// general, any file format is fair game, although it should generally be
/// limited to reasonably popular open formats. For other cases, you can add
/// types to each invocation of ripgrep with the '--type-add' flag.
///
/// If you would like to add or improve this list, please file a PR:
/// <https://github.com/BurntSushi/ripgrep>.
///
/// Please try to keep this list sorted lexicographically and wrapped to 79
/// columns (inclusive).
#[rustfmt::skip]
pub const DEFAULT_TYPES: &[(&str, &[&str])] = &[
    ("agda", &["*.agda", "*.lagda"]),
    ("aidl", &["*.aidl"]),
    ("amake", &["*.mk", "*.bp"]),
    ("asciidoc", &["*.adoc", "*.asc", "*.asciidoc"]),
    ("asm", &["*.asm", "*.s", "*.S"]),
    ("asp", &[
        "*.aspx", "*.aspx.cs", "*.aspx.vb", "*.ascx", "*.ascx.cs", "*.ascx.vb",
    ]),
    ("ats", &["*.ats", "*.dats", "*.sats", "*.hats"]),
    ("avro", &["*.avdl", "*.avpr", "*.avsc"]),
    ("awk", &["*.awk"]),
    ("bazel", &["*.bazel", "*.bzl", "*.BUILD", "*.bazelrc", "BUILD", "WORKSPACE"]),
    ("bitbake", &["*.bb", "*.bbappend", "*.bbclass", "*.conf", "*.inc"]),
    ("brotli", &["*.br"]),
    ("buildstream", &["*.bst"]),
    ("bzip2", &["*.bz2", "*.tbz2"]),
    ("c", &["*.[chH]", "*.[chH].in", "*.cats"]),
    ("cabal", &["*.cabal"]),
    ("cbor", &["*.cbor"]),
    ("ceylon", &["*.ceylon"]),
    ("clojure", &["*.clj", "*.cljc", "*.cljs", "*.cljx"]),
    ("cmake", &["*.cmake", "CMakeLists.txt"]),
    ("coffeescript", &["*.coffee"]),
    ("config", &["*.cfg", "*.conf", "*.config", "*.ini"]),
    ("coq", &["*.v"]),
    ("cpp", &[
        "*.[ChH]", "*.cc", "*.[ch]pp", "*.[ch]xx", "*.hh",  "*.inl",
        "*.[ChH].in", "*.cc.in", "*.[ch]pp.in", "*.[ch]xx.in", "*.hh.in",
    ]),
    ("creole", &["*.creole"]),
    ("crystal", &["Projectfile", "*.cr", "*.ecr", "shard.yml"]),
    ("cs", &["*.cs"]),
    ("csharp", &["*.cs"]),
    ("cshtml", &["*.cshtml"]),
    ("css", &["*.css", "*.scss"]),
    ("csv", &["*.csv"]),
    ("cuda", &["*.cu", "*.cuh"]),
    ("cython", &["*.pyx", "*.pxi", "*.pxd"]),
    ("d", &["*.d"]),
    ("dart", &["*.dart"]),
    ("dhall", &["*.dhall"]),
    ("diff", &["*.patch", "*.diff"]),
    ("docker", &["*Dockerfile*"]),
    ("dvc", &["Dvcfile", "*.dvc"]),
    ("ebuild", &["*.ebuild"]),
    ("edn", &["*.edn"]),
    ("elisp", &["*.el"]),
    ("elixir", &["*.ex", "*.eex", "*.exs"]),
    ("elm", &["*.elm"]),
    ("erb", &["*.erb"]),
    ("erlang", &["*.erl", "*.hrl"]),
    ("fidl", &["*.fidl"]),
    ("fish", &["*.fish"]),
    ("flatbuffers", &["*.fbs"]),
    ("fortran", &[
        "*.f", "*.F", "*.f77", "*.F77", "*.pfo",
        "*.f90", "*.F90", "*.f95", "*.F95",
    ]),
    ("fsharp", &["*.fs", "*.fsx", "*.fsi"]),
    ("fut", &["*.fut"]),
    ("gap", &["*.g", "*.gap", "*.gi", "*.gd", "*.tst"]),
    ("gn", &["*.gn", "*.gni"]),
    ("go", &["*.go"]),
    ("gradle", &["*.gradle"]),
    ("groovy", &["*.groovy", "*.gradle"]),
    ("gzip", &["*.gz", "*.tgz"]),
    ("h", &["*.h", "*.hpp"]),
    ("haml", &["*.haml"]),
    ("haskell", &["*.hs", "*.lhs", "*.cpphs", "*.c2hs", "*.hsc"]),
    ("hbs", &["*.hbs"]),
    ("hs", &["*.hs", "*.lhs"]),
    ("html", &["*.htm", "*.html", "*.ejs"]),
    ("idris", &["*.idr", "*.lidr"]),
    ("janet", &["*.janet"]),
    ("java", &["*.java", "*.jsp", "*.jspx", "*.properties"]),
    ("jinja", &["*.j2", "*.jinja", "*.jinja2"]),
    ("jl", &["*.jl"]),
    ("js", &["*.js", "*.jsx", "*.vue"]),
    ("json", &["*.json", "composer.lock"]),
    ("jsonl", &["*.jsonl"]),
    ("julia", &["*.jl"]),
    ("jupyter", &["*.ipynb", "*.jpynb"]),
    ("k", &["*.k"]),
    ("kotlin", &["*.kt", "*.kts"]),
    ("less", &["*.less"]),
    ("license", &[
        // General
        "COPYING", "COPYING[.-]*",
        "COPYRIGHT", "COPYRIGHT[.-]*",
        "EULA", "EULA[.-]*",
        "licen[cs]e", "licen[cs]e.*",
        "LICEN[CS]E", "LICEN[CS]E[.-]*", "*[.-]LICEN[CS]E*",
        "NOTICE", "NOTICE[.-]*",
        "PATENTS", "PATENTS[.-]*",
        "UNLICEN[CS]E", "UNLICEN[CS]E[.-]*",
        // GPL (gpl.txt, etc.)
        "agpl[.-]*",
        "gpl[.-]*",
        "lgpl[.-]*",
        // Other license-specific (APACHE-2.0.txt, etc.)
        "AGPL-*[0-9]*",
        "APACHE-*[0-9]*",
        "BSD-*[0-9]*",
        "CC-BY-*",
        "GFDL-*[0-9]*",
        "GNU-*[0-9]*",
        "GPL-*[0-9]*",
        "LGPL-*[0-9]*",
        "MIT-*[0-9]*",
        "MPL-*[0-9]*",
        "OFL-*[0-9]*",
    ]),
    ("lisp", &["*.el", "*.jl", "*.lisp", "*.lsp", "*.sc", "*.scm"]),
    ("lock", &["*.lock", "package-lock.json"]),
    ("log", &["*.log"]),
    ("lua", &["*.lua"]),
    ("lz4", &["*.lz4"]),
    ("lzma", &["*.lzma"]),
    ("m4", &["*.ac", "*.m4"]),
    ("make", &[
        "[Gg][Nn][Uu]makefile", "[Mm]akefile",
        "[Gg][Nn][Uu]makefile.am", "[Mm]akefile.am",
        "[Gg][Nn][Uu]makefile.in", "[Mm]akefile.in",
        "*.mk", "*.mak"
    ]),
    ("mako", &["*.mako", "*.mao"]),
    ("man", &["*.[0-9lnpx]", "*.[0-9][cEFMmpSx]"]),
    ("markdown", &["*.markdown", "*.md", "*.mdown", "*.mkdn"]),
    ("matlab", &["*.m"]),
    ("md", &["*.markdown", "*.md", "*.mdown", "*.mkdn"]),
    ("meson", &["meson.build", "meson_options.txt"]),
    ("minified", &["*.min.html", "*.min.css", "*.min.js"]),
    ("mint", &["*.mint"]),
    ("mk", &["mkfile"]),
    ("ml", &["*.ml"]),
    ("msbuild", &[
        "*.csproj", "*.fsproj", "*.vcxproj", "*.proj", "*.props", "*.targets",
    ]),
    ("nim", &["*.nim", "*.nimf", "*.nimble", "*.nims"]),
    ("nix", &["*.nix"]),
    ("objc", &["*.h", "*.m"]),
    ("objcpp", &["*.h", "*.mm"]),
    ("ocaml", &["*.ml", "*.mli", "*.mll", "*.mly"]),
    ("org", &["*.org", "*.org_archive"]),
    ("pascal", &["*.pas", "*.dpr", "*.lpr", "*.pp", "*.inc"]),
    ("pdf", &["*.pdf"]),
    ("perl", &["*.perl", "*.pl", "*.PL", "*.plh", "*.plx", "*.pm", "*.t"]),
    ("php", &["*.php", "*.php3", "*.php4", "*.php5", "*.phtml"]),
    ("po", &["*.po"]),
    ("pod", &["*.pod"]),
    ("postscript", &["*.eps", "*.ps"]),
    ("protobuf", &["*.proto"]),
    ("ps", &["*.cdxml", "*.ps1", "*.ps1xml", "*.psd1", "*.psm1"]),
    ("puppet", &["*.erb", "*.pp", "*.rb"]),
    ("purs", &["*.purs"]),
    ("py", &["*.py"]),
    ("qmake", &["*.pro", "*.pri", "*.prf"]),
    ("qml", &["*.qml"]),
    ("r", &["*.R", "*.r", "*.Rmd", "*.Rnw"]),
    ("racket", &["*.rkt"]),
    ("rdoc", &["*.rdoc"]),
    ("readme", &["README*", "*README"]),
    ("red", &["*.r", "*.red", "*.reds"]),
    ("robot", &["*.robot"]),
    ("rst", &["*.rst"]),
    ("ruby", &[
        // Idiomatic files
        "config.ru", "Gemfile", ".irbrc", "Rakefile",
        // Extensions
        "*.gemspec", "*.rb", "*.rbw"
    ]),
    ("rust", &["*.rs"]),
    ("sass", &["*.sass", "*.scss"]),
    ("scala", &["*.scala", "*.sbt"]),
    ("sh", &[
        // Portable/misc. init files
        ".login", ".logout", ".profile", "profile",
        // bash-specific init files
        ".bash_login", "bash_login",
        ".bash_logout", "bash_logout",
        ".bash_profile", "bash_profile",
        ".bashrc", "bashrc", "*.bashrc",
        // csh-specific init files
        ".cshrc", "*.cshrc",
        // ksh-specific init files
        ".kshrc", "*.kshrc",
        // tcsh-specific init files
        ".tcshrc",
        // zsh-specific init files
        ".zshenv", "zshenv",
        ".zlogin", "zlogin",
        ".zlogout", "zlogout",
        ".zprofile", "zprofile",
        ".zshrc", "zshrc",
        // Extensions
        "*.bash", "*.csh", "*.ksh", "*.sh", "*.tcsh", "*.zsh",
    ]),
    ("slim", &["*.skim", "*.slim", "*.slime"]),
    ("smarty", &["*.tpl"]),
    ("sml", &["*.sml", "*.sig"]),
    ("soy", &["*.soy"]),
    ("spark", &["*.spark"]),
    ("spec", &["*.spec"]),
    ("sql", &["*.sql", "*.psql"]),
    ("stylus", &["*.styl"]),
    ("sv", &["*.v", "*.vg", "*.sv", "*.svh", "*.h"]),
    ("svg", &["*.svg"]),
    ("swift", &["*.swift"]),
    ("swig", &["*.def", "*.i"]),
    ("systemd", &[
        "*.automount", "*.conf", "*.device", "*.link", "*.mount", "*.path",
        "*.scope", "*.service", "*.slice", "*.socket", "*.swap", "*.target",
        "*.timer",
    ]),
    ("taskpaper", &["*.taskpaper"]),
    ("tcl", &["*.tcl"]),
    ("tex", &["*.tex", "*.ltx", "*.cls", "*.sty", "*.bib", "*.dtx", "*.ins"]),
    ("texinfo", &["*.texi"]),
    ("textile", &["*.textile"]),
    ("tf", &["*.tf"]),
    ("thrift", &["*.thrift"]),
    ("toml", &["*.toml", "Cargo.lock"]),
    ("ts", &["*.ts", "*.tsx"]),
    ("twig", &["*.twig"]),
    ("txt", &["*.txt"]),
    ("typoscript", &["*.typoscript", "*.ts"]),
    ("vala", &["*.vala"]),
    ("vb", &["*.vb"]),
    ("vcl", &["*.vcl"]),
    ("verilog", &["*.v", "*.vh", "*.sv", "*.svh"]),
    ("vhdl", &["*.vhd", "*.vhdl"]),
    ("vim", &["*.vim"]),
    ("vimscript", &["*.vim"]),
    ("webidl", &["*.idl", "*.webidl", "*.widl"]),
    ("wiki", &["*.mediawiki", "*.wiki"]),
    ("xml", &[
        "*.xml", "*.xml.dist", "*.dtd", "*.xsl", "*.xslt", "*.xsd", "*.xjb",
        "*.rng", "*.sch", "*.xhtml",
    ]),
    ("xz", &["*.xz", "*.txz"]),
    ("yacc", &["*.y"]),
    ("yaml", &["*.yaml", "*.yml"]),
    ("yang", &["*.yang"]),
    ("z", &["*.Z"]),
    ("zig", &["*.zig"]),
    ("zsh", &[
        ".zshenv", "zshenv",
        ".zlogin", "zlogin",
        ".zlogout", "zlogout",
        ".zprofile", "zprofile",
        ".zshrc", "zshrc",
        "*.zsh",
    ]),
    ("zstd", &["*.zst", "*.zstd"]),
];
<|endoftext|>


<commit_before>// Copyright:: Copyright (c) 2015-2016 The Habitat Maintainers
//
// The terms of the Evaluation Agreement (Habitat) between Chef Software Inc.
// and the party accessing this file ("Licensee") apply to Licensee's use of
// the Software until such time that the Software is made available under an
// open source license such as the Apache 2.0 License.

extern crate habitat_core as hcore;
extern crate habitat_depot as depot;
extern crate habitat_depot_core as depot_core;
#[macro_use]
extern crate clap;
extern crate env_logger;
#[macro_use]
extern crate log;

use std::env;
use std::path::PathBuf;
use std::process;
use std::str::FromStr;

use depot::{server, Config, Error, Result};
use depot::data_store::{self, Cursor, Database, Transaction};
use depot_core::data_object;
use hcore::fs;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

fn main() {
    env_logger::init().unwrap();
    let matches = app().get_matches();
    debug!("CLI matches: {:?}", matches);
    let config = match config_from_args(&matches) {
        Ok(result) => result,
        Err(e) => return exit_with(e, 1),
    };
    match dispatch(config, &matches) {
        Ok(_) => std::process::exit(0),
        Err(e) => exit_with(e, 1),
    }
}

fn app<'a, 'b>() -> clap::App<'a, 'b> {
    clap_app!(BldrDepot =>
        (version: VERSION)
        (about: "Manage a package Depot")
        (@setting VersionlessSubcommands)
        (@setting SubcommandRequiredElseHelp)
        (@arg path: -p --path +takes_value "Filepath to service storage for the Depot service")
        (@subcommand start =>
            (about: "Run a bldr package Depot")
            (@arg port: --port +takes_value "Listen port. [default: 9632]")
        )
        (@subcommand repair =>
            (about: "Verify and repair data integrity of the package Depot")
        )
        (@subcommand view =>
            (@subcommand create =>
                (about: "Create a new view over the package Depot")
                (@arg view: <view> +required "Name of the view to create")
            )
            (@subcommand list =>
                (about: "List views in the package Depot")
            )
        )
    )
}

fn config_from_args(matches: &clap::ArgMatches) -> Result<Config> {
    let cmd = matches.subcommand_name().unwrap();
    let args = matches.subcommand_matches(cmd).unwrap();
    let mut config = Config::new();
    if let Some(port) = args.value_of("port") {
        if let Some(port) = u16::from_str(port).ok() {
            config.port = depot::ListenPort(port);
        } else {
            return Err(Error::BadPort(port.to_string()));
        }
    }
    config.path = args.value_of("path").unwrap_or(&default_path()).to_string();
    Ok(config)
}

fn dispatch(config: Config, matches: &clap::ArgMatches) -> Result<()> {
    match matches.subcommand_name() {
        Some("start") => start(&config),
        Some("repair") => repair(&config),
        Some(cmd @ "view") => {
            let args = matches.subcommand_matches(cmd).unwrap();
            match args.subcommand_name() {
                Some(cmd @ "create") => {
                    let args = args.subcommand_matches(cmd).unwrap();
                    let name = args.value_of("view").unwrap();
                    view_create(name, &config)
                }
                Some("list") => view_list(&config),
                Some(cmd) => {
                    debug!("Dispatch failed, no match for command: {:?}", cmd);
                    Ok(())
                }
                None => Ok(()),
            }
        }
        Some(cmd) => {
            debug!("Dispatch failed, no match for command: {:?}", cmd);
            Ok(())
        }
        None => Ok(()),
    }
}

/// Starts the depot server.
///
/// # Failures
///
/// * Fails if the depot server fails to start - canot bind to the port, etc.
fn start(config: &Config) -> Result<()> {
    println!("Starting package Depot at {}", &config.path);
    println!("Depot listening on {:?}", config.depot_addr());
    server::run(&config)
}

/// Analyzes the integrity of the depot's metadata by comparing the metadata with the packages
/// on disk. If a package is found on disk that is not present in the metadata it is added to the
/// metadata and if an entry in the metadata doesn't have a matching package archive on disk the
/// entry is dropped from the database.
///
/// # Failures
///
/// * The database cannot be read
/// * A write transaction cannot be acquired
pub fn repair(config: &Config) -> Result<()> {
    // JW TODO: should pass config to depot, not this path
    let depot = try!(depot::Depot::new(config.path.clone()));
    let report = try!(depot::doctor::repair(&depot));
    println!("Report: {:?}", &report);
    Ok(())
}

/// Create a view with the given name in the depot.
///
/// # Failures
///
/// * The database cannot be read
/// * A write transaction cannot be acquired.
fn view_create(name: &str, config: &Config) -> Result<()> {
    // JW TODO: should pass config to depot, not this path
    let depot = try!(depot::Depot::new(config.path.clone()));
    let txn = try!(depot.datastore.views.txn_rw());
    let object = data_object::View::new(name);
    try!(depot.datastore.views.write(&txn, &object));
    Ok(())
}

/// List all views in the database.
///
/// # Failures
///
/// * The database cannot be read
/// * A read transaction cannot be acquired.
fn view_list(config: &Config) -> Result<()> {
    // JW TODO: should pass config to depot, not this path
    let depot = try!(depot::Depot::new(config.path.clone()));
    let mut views: Vec<data_object::View> = vec![];
    let txn = try!(depot.datastore.views.txn_ro());
    let mut cursor = try!(txn.cursor_ro());
    match cursor.first() {
        Err(Error::MdbError(data_store::MdbError::NotFound)) => {
            println!("No views. Create one with `hab-depot view create`.");
            return Ok(());
        }
        Err(e) => return Err(e),
        Ok((_, value)) => views.push(value),
    }
    loop {
        match cursor.next() {
            Ok((_, value)) => views.push(value),
            Err(_) => break,
        }
    }
    println!("Listing {} views", views.len());
    for view in views.iter() {
        println!("     {}", view);
    }
    Ok(())
}

fn default_path() -> String {
    let arg0 = env::args().next().map(|p| PathBuf::from(p));
    let progname = arg0.as_ref().and_then(|p| p.file_stem()).and_then(|p| p.to_str()).unwrap();
    fs::svc_path(progname).join("data").to_string_lossy().into_owned()
}

fn exit_with(err: Error, code: i32) {
    println!("{:?}", err);
    process::exit(code)
}
<commit_msg>String updates to hab-depot command help<commit_after>// Copyright:: Copyright (c) 2015-2016 The Habitat Maintainers
//
// The terms of the Evaluation Agreement (Habitat) between Chef Software Inc.
// and the party accessing this file ("Licensee") apply to Licensee's use of
// the Software until such time that the Software is made available under an
// open source license such as the Apache 2.0 License.

extern crate habitat_core as hcore;
extern crate habitat_depot as depot;
extern crate habitat_depot_core as depot_core;
#[macro_use]
extern crate clap;
extern crate env_logger;
#[macro_use]
extern crate log;

use std::env;
use std::path::PathBuf;
use std::process;
use std::str::FromStr;

use depot::{server, Config, Error, Result};
use depot::data_store::{self, Cursor, Database, Transaction};
use depot_core::data_object;
use hcore::fs;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

fn main() {
    env_logger::init().unwrap();
    let matches = app().get_matches();
    debug!("CLI matches: {:?}", matches);
    let config = match config_from_args(&matches) {
        Ok(result) => result,
        Err(e) => return exit_with(e, 1),
    };
    match dispatch(config, &matches) {
        Ok(_) => std::process::exit(0),
        Err(e) => exit_with(e, 1),
    }
}

fn app<'a, 'b>() -> clap::App<'a, 'b> {
    clap_app!(BldrDepot =>
        (version: VERSION)
        (about: "Manage a package Depot")
        (@setting VersionlessSubcommands)
        (@setting SubcommandRequiredElseHelp)
        (@arg path: -p --path +takes_value "Filepath to service storage for the Depot service")
        (@subcommand start =>
            (about: "Run a Habitat package Depot")
            (@arg port: --port +takes_value "Listen port. [default: 9632]")
        )
        (@subcommand repair =>
            (about: "Verify and repair data integrity of the package Depot")
        )
        (@subcommand view =>
            (about: "Creates or lists views in the package Depot")
            (@subcommand create =>
                (about: "Create a new view over the package Depot")
                (@arg view: <view> +required "Name of the view to create")
            )
            (@subcommand list =>
                (about: "List views in the package Depot")
            )
        )
    )
}

fn config_from_args(matches: &clap::ArgMatches) -> Result<Config> {
    let cmd = matches.subcommand_name().unwrap();
    let args = matches.subcommand_matches(cmd).unwrap();
    let mut config = Config::new();
    if let Some(port) = args.value_of("port") {
        if let Some(port) = u16::from_str(port).ok() {
            config.port = depot::ListenPort(port);
        } else {
            return Err(Error::BadPort(port.to_string()));
        }
    }
    config.path = args.value_of("path").unwrap_or(&default_path()).to_string();
    Ok(config)
}

fn dispatch(config: Config, matches: &clap::ArgMatches) -> Result<()> {
    match matches.subcommand_name() {
        Some("start") => start(&config),
        Some("repair") => repair(&config),
        Some(cmd @ "view") => {
            let args = matches.subcommand_matches(cmd).unwrap();
            match args.subcommand_name() {
                Some(cmd @ "create") => {
                    let args = args.subcommand_matches(cmd).unwrap();
                    let name = args.value_of("view").unwrap();
                    view_create(name, &config)
                }
                Some("list") => view_list(&config),
                Some(cmd) => {
                    debug!("Dispatch failed, no match for command: {:?}", cmd);
                    Ok(())
                }
                None => Ok(()),
            }
        }
        Some(cmd) => {
            debug!("Dispatch failed, no match for command: {:?}", cmd);
            Ok(())
        }
        None => Ok(()),
    }
}

/// Starts the depot server.
///
/// # Failures
///
/// * Fails if the depot server fails to start - canot bind to the port, etc.
fn start(config: &Config) -> Result<()> {
    println!("Starting package Depot at {}", &config.path);
    println!("Depot listening on {:?}", config.depot_addr());
    server::run(&config)
}

/// Analyzes the integrity of the depot's metadata by comparing the metadata with the packages
/// on disk. If a package is found on disk that is not present in the metadata it is added to the
/// metadata and if an entry in the metadata doesn't have a matching package archive on disk the
/// entry is dropped from the database.
///
/// # Failures
///
/// * The database cannot be read
/// * A write transaction cannot be acquired
pub fn repair(config: &Config) -> Result<()> {
    // JW TODO: should pass config to depot, not this path
    let depot = try!(depot::Depot::new(config.path.clone()));
    let report = try!(depot::doctor::repair(&depot));
    println!("Report: {:?}", &report);
    Ok(())
}

/// Create a view with the given name in the depot.
///
/// # Failures
///
/// * The database cannot be read
/// * A write transaction cannot be acquired.
fn view_create(name: &str, config: &Config) -> Result<()> {
    // JW TODO: should pass config to depot, not this path
    let depot = try!(depot::Depot::new(config.path.clone()));
    let txn = try!(depot.datastore.views.txn_rw());
    let object = data_object::View::new(name);
    try!(depot.datastore.views.write(&txn, &object));
    Ok(())
}

/// List all views in the database.
///
/// # Failures
///
/// * The database cannot be read
/// * A read transaction cannot be acquired.
fn view_list(config: &Config) -> Result<()> {
    // JW TODO: should pass config to depot, not this path
    let depot = try!(depot::Depot::new(config.path.clone()));
    let mut views: Vec<data_object::View> = vec![];
    let txn = try!(depot.datastore.views.txn_ro());
    let mut cursor = try!(txn.cursor_ro());
    match cursor.first() {
        Err(Error::MdbError(data_store::MdbError::NotFound)) => {
            println!("No views. Create one with `hab-depot view create`.");
            return Ok(());
        }
        Err(e) => return Err(e),
        Ok((_, value)) => views.push(value),
    }
    loop {
        match cursor.next() {
            Ok((_, value)) => views.push(value),
            Err(_) => break,
        }
    }
    println!("Listing {} views", views.len());
    for view in views.iter() {
        println!("     {}", view);
    }
    Ok(())
}

fn default_path() -> String {
    let arg0 = env::args().next().map(|p| PathBuf::from(p));
    let progname = arg0.as_ref().and_then(|p| p.file_stem()).and_then(|p| p.to_str()).unwrap();
    fs::svc_path(progname).join("data").to_string_lossy().into_owned()
}

fn exit_with(err: Error, code: i32) {
    println!("{:?}", err);
    process::exit(code)
}
<|endoftext|>


<commit_before>use rustup::Notification;
use rustup_dist::Notification as In;
use rustup_utils::tty;
use rustup_utils::Notification as Un;
use std::collections::VecDeque;
use std::fmt;
use term;
use time::precise_time_s;

/// Keep track of this many past download amounts
const DOWNLOAD_TRACK_COUNT: usize = 5;

/// Tracks download progress and displays information about it to a terminal.
pub struct DownloadTracker {
    /// Content-Length of the to-be downloaded object.
    content_len: Option<u64>,
    /// Total data downloaded in bytes.
    total_downloaded: usize,
    /// Data downloaded this second.
    downloaded_this_sec: usize,
    /// Keeps track of amount of data downloaded every last few secs.
    /// Used for averaging the download speed.
    downloaded_last_few_secs: VecDeque<usize>,
    /// Time stamp of the last second
    last_sec: Option<f64>,
    /// How many seconds have elapsed since the download started
    seconds_elapsed: u32,
    /// The terminal we write the information to.
    term: Option<Box<term::StdoutTerminal>>,
    /// Whether we displayed progress for the download or not.
    ///
    /// If the download is quick enough, we don't have time to
    /// display the progress info.
    /// In that case, we do not want to do some cleanup stuff we normally do.
    displayed_progress: bool,
}

impl DownloadTracker {
    /// Creates a new DownloadTracker.
    pub fn new() -> Self {
        DownloadTracker {
            content_len: None,
            total_downloaded: 0,
            downloaded_this_sec: 0,
            downloaded_last_few_secs: VecDeque::with_capacity(DOWNLOAD_TRACK_COUNT),
            seconds_elapsed: 0,
            last_sec: None,
            term: term::stdout(),
            displayed_progress: false,
        }
    }

    pub fn handle_notification(&mut self, n: &Notification) -> bool {
        match *n {
            Notification::Install(In::Utils(Un::DownloadContentLengthReceived(content_len))) => {
                self.content_length_received(content_len);

                true
            }
            Notification::Install(In::Utils(Un::DownloadDataReceived(data))) => {
                if tty::stdout_isatty() && self.term.is_some() {
                    self.data_received(data.len());
                }
                true
            }
            Notification::Install(In::Utils(Un::DownloadFinished)) => {
                self.download_finished();
                true
            }
            _ => false,
        }
    }

    /// Notifies self that Content-Length information has been received.
    pub fn content_length_received(&mut self, content_len: u64) {
        self.content_len = Some(content_len);
    }
    /// Notifies self that data of size `len` has been received.
    pub fn data_received(&mut self, len: usize) {
        self.total_downloaded += len;
        self.downloaded_this_sec += len;

        let current_time = precise_time_s();

        match self.last_sec {
            None => self.last_sec = Some(current_time),
            Some(start) => {
                let elapsed = current_time - start;
                if elapsed >= 1.0 {
                    self.seconds_elapsed += 1;

                    self.display();
                    self.last_sec = Some(current_time);
                    if self.downloaded_last_few_secs.len() == DOWNLOAD_TRACK_COUNT {
                        self.downloaded_last_few_secs.pop_back();
                    }
                    self.downloaded_last_few_secs
                        .push_front(self.downloaded_this_sec);
                    self.downloaded_this_sec = 0;
                }
            }
        }
    }
    /// Notifies self that the download has finished.
    pub fn download_finished(&mut self) {
        if self.displayed_progress {
            // Display the finished state
            self.display();
            let _ = writeln!(self.term.as_mut().unwrap(), "");
        }
        self.prepare_for_new_download();
    }
    /// Resets the state to be ready for a new download.
    fn prepare_for_new_download(&mut self) {
        self.content_len = None;
        self.total_downloaded = 0;
        self.downloaded_this_sec = 0;
        self.downloaded_last_few_secs.clear();
        self.seconds_elapsed = 0;
        self.last_sec = None;
        self.displayed_progress = false;
    }
    /// Display the tracked download information to the terminal.
    fn display(&mut self) {
        let total_h = HumanReadable(self.total_downloaded as f64);
        let sum = self
            .downloaded_last_few_secs
            .iter()
            .fold(0., |a, &v| a + v as f64);
        let len = self.downloaded_last_few_secs.len();
        let speed = if len > 0 { sum / len as f64 } else { 0. };
        let speed_h = HumanReadable(speed);

        match self.content_len {
            Some(content_len) => {
                let content_len = content_len as f64;
                let percent = (self.total_downloaded as f64 / content_len) * 100.;
                let content_len_h = HumanReadable(content_len);
                let remaining = content_len - self.total_downloaded as f64;
                let eta_h = HumanReadable(remaining / speed);
                let _ = write!(
                    self.term.as_mut().unwrap(),
                    "{} / {} ({:3.0} %) {}/s ETA: {:#}",
                    total_h,
                    content_len_h,
                    percent,
                    speed_h,
                    eta_h
                );
            }
            None => {
                let _ = write!(
                    self.term.as_mut().unwrap(),
                    "Total: {} Speed: {}/s",
                    total_h,
                    speed_h
                );
            }
        }
        // delete_line() doesn't seem to clear the line properly.
        // Instead, let's just print some whitespace to clear it.
        let _ = write!(self.term.as_mut().unwrap(), "                ");
        let _ = self.term.as_mut().unwrap().flush();
        let _ = self.term.as_mut().unwrap().carriage_return();
        self.displayed_progress = true;
    }
}

/// Human readable representation of data size in bytes
struct HumanReadable(f64);

impl fmt::Display for HumanReadable {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if f.alternate() {
            // repurposing the alternate mode for ETA
            let sec = self.0;

            if sec.is_infinite() {
                write!(f, "Unknown")
            } else if sec > 48. * 3600. {
                let sec = self.0 as u64;
                let d = sec / (24. * 3600.);
                let h = sec % (24. * 3600.);
                let min = sec % 3600.;
                let sec = sec % 60.;

                write!(f, "{:3} days {:2} h {:2} min {:2} s", d, h, min, sec) // XYZ days PQ h RS min TU s
            } else if sec > 6_000. {
                let sec = self.0 as u64;
                let h = sec / 3600.;
                let min = sec % 3600.;
                let sec = sec % 60.;

                write!(f, "{:3} h {:2} min {:2} s", h, min, sec) // XYZ h PQ min RS s
            } else if sec > 100 {
                let sec = self.0 as u64;
                let min = sec / 60.;
                let sec = sec % 60.;

                write!(f, "{:3} min {:2} s", min, sec) // XYZ min PQ s
            } else {
                write!(f, "{:3.0} s", self.0) // XYZ s
            }
        } else {
            const KIB: f64 = 1024.0;
            const MIB: f64 = KIB * KIB;
            let size = self.0;

            if size >= MIB {
                write!(f, "{:5.1} MiB", size / MIB) // XYZ.P MiB
            } else if size >= KIB {
                write!(f, "{:5.1} KiB", size / KIB)
            } else {
                write!(f, "{:3.0} B", size)
            }
        }
    }
}
<commit_msg>Fix HumanReadable#fmt<commit_after>use rustup::Notification;
use rustup_dist::Notification as In;
use rustup_utils::tty;
use rustup_utils::Notification as Un;
use std::collections::VecDeque;
use std::fmt;
use term;
use time::precise_time_s;

/// Keep track of this many past download amounts
const DOWNLOAD_TRACK_COUNT: usize = 5;

/// Tracks download progress and displays information about it to a terminal.
pub struct DownloadTracker {
    /// Content-Length of the to-be downloaded object.
    content_len: Option<u64>,
    /// Total data downloaded in bytes.
    total_downloaded: usize,
    /// Data downloaded this second.
    downloaded_this_sec: usize,
    /// Keeps track of amount of data downloaded every last few secs.
    /// Used for averaging the download speed.
    downloaded_last_few_secs: VecDeque<usize>,
    /// Time stamp of the last second
    last_sec: Option<f64>,
    /// How many seconds have elapsed since the download started
    seconds_elapsed: u32,
    /// The terminal we write the information to.
    term: Option<Box<term::StdoutTerminal>>,
    /// Whether we displayed progress for the download or not.
    ///
    /// If the download is quick enough, we don't have time to
    /// display the progress info.
    /// In that case, we do not want to do some cleanup stuff we normally do.
    displayed_progress: bool,
}

impl DownloadTracker {
    /// Creates a new DownloadTracker.
    pub fn new() -> Self {
        DownloadTracker {
            content_len: None,
            total_downloaded: 0,
            downloaded_this_sec: 0,
            downloaded_last_few_secs: VecDeque::with_capacity(DOWNLOAD_TRACK_COUNT),
            seconds_elapsed: 0,
            last_sec: None,
            term: term::stdout(),
            displayed_progress: false,
        }
    }

    pub fn handle_notification(&mut self, n: &Notification) -> bool {
        match *n {
            Notification::Install(In::Utils(Un::DownloadContentLengthReceived(content_len))) => {
                self.content_length_received(content_len);

                true
            }
            Notification::Install(In::Utils(Un::DownloadDataReceived(data))) => {
                if tty::stdout_isatty() && self.term.is_some() {
                    self.data_received(data.len());
                }
                true
            }
            Notification::Install(In::Utils(Un::DownloadFinished)) => {
                self.download_finished();
                true
            }
            _ => false,
        }
    }

    /// Notifies self that Content-Length information has been received.
    pub fn content_length_received(&mut self, content_len: u64) {
        self.content_len = Some(content_len);
    }
    /// Notifies self that data of size `len` has been received.
    pub fn data_received(&mut self, len: usize) {
        self.total_downloaded += len;
        self.downloaded_this_sec += len;

        let current_time = precise_time_s();

        match self.last_sec {
            None => self.last_sec = Some(current_time),
            Some(start) => {
                let elapsed = current_time - start;
                if elapsed >= 1.0 {
                    self.seconds_elapsed += 1;

                    self.display();
                    self.last_sec = Some(current_time);
                    if self.downloaded_last_few_secs.len() == DOWNLOAD_TRACK_COUNT {
                        self.downloaded_last_few_secs.pop_back();
                    }
                    self.downloaded_last_few_secs
                        .push_front(self.downloaded_this_sec);
                    self.downloaded_this_sec = 0;
                }
            }
        }
    }
    /// Notifies self that the download has finished.
    pub fn download_finished(&mut self) {
        if self.displayed_progress {
            // Display the finished state
            self.display();
            let _ = writeln!(self.term.as_mut().unwrap(), "");
        }
        self.prepare_for_new_download();
    }
    /// Resets the state to be ready for a new download.
    fn prepare_for_new_download(&mut self) {
        self.content_len = None;
        self.total_downloaded = 0;
        self.downloaded_this_sec = 0;
        self.downloaded_last_few_secs.clear();
        self.seconds_elapsed = 0;
        self.last_sec = None;
        self.displayed_progress = false;
    }
    /// Display the tracked download information to the terminal.
    fn display(&mut self) {
        let total_h = HumanReadable(self.total_downloaded as f64);
        let sum = self
            .downloaded_last_few_secs
            .iter()
            .fold(0., |a, &v| a + v as f64);
        let len = self.downloaded_last_few_secs.len();
        let speed = if len > 0 { sum / len as f64 } else { 0. };
        let speed_h = HumanReadable(speed);

        match self.content_len {
            Some(content_len) => {
                let content_len = content_len as f64;
                let percent = (self.total_downloaded as f64 / content_len) * 100.;
                let content_len_h = HumanReadable(content_len);
                let remaining = content_len - self.total_downloaded as f64;
                let eta_h = HumanReadable(remaining / speed);
                let _ = write!(
                    self.term.as_mut().unwrap(),
                    "{} / {} ({:3.0} %) {}/s ETA: {:#}",
                    total_h,
                    content_len_h,
                    percent,
                    speed_h,
                    eta_h
                );
            }
            None => {
                let _ = write!(
                    self.term.as_mut().unwrap(),
                    "Total: {} Speed: {}/s",
                    total_h,
                    speed_h
                );
            }
        }
        // delete_line() doesn't seem to clear the line properly.
        // Instead, let's just print some whitespace to clear it.
        let _ = write!(self.term.as_mut().unwrap(), "                ");
        let _ = self.term.as_mut().unwrap().flush();
        let _ = self.term.as_mut().unwrap().carriage_return();
        self.displayed_progress = true;
    }
}

/// Human readable representation of data size in bytes
struct HumanReadable(f64);

impl fmt::Display for HumanReadable {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if f.alternate() {
            // repurposing the alternate mode for ETA
            let sec = self.0;

            if sec.is_infinite() {
                write!(f, "Unknown")
            } else if sec > 48. * 3600. {
                let sec = self.0;
                let d = sec / (24. * 3600.);
                let h = sec % (24. * 3600.);
                let min = sec % 3600.;
                let sec = sec % 60.;

                write!(f, "{:3} days {:2} h {:2} min {:2} s", d, h, min, sec) // XYZ days PQ h RS min TU s
            } else if sec > 6_000. {
                let sec = self.0;
                let h = sec / 3600.;
                let min = sec % 3600.;
                let sec = sec % 60.;

                write!(f, "{:3} h {:2} min {:2} s", h, min, sec) // XYZ h PQ min RS s
            } else if sec > 100. {
                let sec = self.0;
                let min = sec / 60.;
                let sec = sec % 60.;

                write!(f, "{:3} min {:2} s", min, sec) // XYZ min PQ s
            } else {
                write!(f, "{:3.0} s", self.0) // XYZ s
            }
        } else {
            const KIB: f64 = 1024.0;
            const MIB: f64 = KIB * KIB;
            let size = self.0;

            if size >= MIB {
                write!(f, "{:5.1} MiB", size / MIB) // XYZ.P MiB
            } else if size >= KIB {
                write!(f, "{:5.1} KiB", size / KIB)
            } else {
                write!(f, "{:3.0} B", size)
            }
        }
    }
}
<|endoftext|>


<commit_before>use byteorder::{ByteOrder, BigEndian};
use rocksdb;

use {RocksDBIndexStore, RocksDBIndexReader};
use segment::Segment;
use key_builder::KeyBuilder;

#[derive(Debug)]
pub struct SegmentStatistics {
    total_docs: i64,
    deleted_docs: i64,
}

impl SegmentStatistics {
    fn read(reader: &RocksDBIndexReader, segment: u32) -> Result<SegmentStatistics, rocksdb::Error> {
        // Total docs
        let kb = KeyBuilder::segment_stat(segment, b"total_docs");

        let total_docs = match try!(reader.snapshot.get(&kb.key())) {
            Some(val_bytes) => {
                BigEndian::read_i64(&val_bytes)
            }
            None => 0,
        };

        // Deleted docs
        let kb = KeyBuilder::segment_stat(segment, b"deleted_docs");

        let deleted_docs = match try!(reader.snapshot.get(&kb.key())) {
            Some(val_bytes) => {
                BigEndian::read_i64(&val_bytes)
            }
            None => 0,
        };

        Ok(SegmentStatistics {
            total_docs: total_docs,
            deleted_docs: deleted_docs,
        })
    }

    #[inline]
    pub fn total_docs(&self) -> i64 {
        self.total_docs
    }

    #[inline]
    pub fn deleted_docs(&self) -> i64 {
        self.deleted_docs
    }
}

impl RocksDBIndexStore {
    pub fn get_segment_statistics(&self) -> Result<Vec<(u32, SegmentStatistics)>, rocksdb::Error> {
        let mut segment_stats = Vec::new();
        let reader = self.reader();

        for segment in self.segments.iter_active(&reader) {
            let stats = try!(SegmentStatistics::read(&reader, segment.id()));
            segment_stats.push((segment.id(), stats));
        }

        Ok(segment_stats)
    }
}
<commit_msg>Make SegmentStatistics generic<commit_after>use byteorder::{ByteOrder, BigEndian};
use rocksdb;

use {RocksDBIndexStore, RocksDBIndexReader};
use segment::Segment;
use key_builder::KeyBuilder;

#[derive(Debug)]
pub struct SegmentStatistics {
    total_docs: i64,
    deleted_docs: i64,
}

impl SegmentStatistics {
    fn read<S: Segment>(reader: &RocksDBIndexReader, segment: &S) -> Result<SegmentStatistics, rocksdb::Error> {
        let total_docs = try!(segment.load_statistic(b"total_docs")).unwrap_or(0);
        let deleted_docs = try!(segment.load_statistic(b"deleted_docs")).unwrap_or(0);

        Ok(SegmentStatistics {
            total_docs: total_docs,
            deleted_docs: deleted_docs,
        })
    }

    #[inline]
    pub fn total_docs(&self) -> i64 {
        self.total_docs
    }

    #[inline]
    pub fn deleted_docs(&self) -> i64 {
        self.deleted_docs
    }
}

impl RocksDBIndexStore {
    pub fn get_segment_statistics(&self) -> Result<Vec<(u32, SegmentStatistics)>, rocksdb::Error> {
        let mut segment_stats = Vec::new();
        let reader = self.reader();

        for segment in self.segments.iter_active(&reader) {
            let stats = try!(SegmentStatistics::read(&reader, &segment));
            segment_stats.push((segment.id(), stats));
        }

        Ok(segment_stats)
    }
}
<|endoftext|>


<commit_before>#![allow(non_snake_case)]
// #![cfg_attr(test, feature(test))]
extern crate time;
extern crate num_cpus;
extern crate scoped_threadpool;

use time::precise_time_s;
use scoped_threadpool::Pool;
use std::mem;

///////////////////////////////////////////////////////////////////////////////
// 1. Calibration
///////////////////////////////////////////////////////////////////////////////
const ALPHA: f64 = 0.33333333333;  // Elasticity of output w.r.t. capital
const BETA: f64 = 0.95;           // Discount factor

// Productivity values
const PRODUCTIVITY: [f64; 5] = [0.9792, 0.9896, 1.0000, 1.0106, 1.0212];

// Transition matrix
static TRANSITIONS: [[f64; 5]; 5] = [[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],
                                     [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],
                                     [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],
                                     [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],
                                     [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]];

// Dimensions to generate the grid of capital
const GRID_CAPITAL: usize = 17820;
const GRID_PRODUCTIVITY: usize = 5;

fn solve(print: bool) -> f64 {
  /////////////////////////////////////////////////////////////////////////////
  // 2. Steady State
  /////////////////////////////////////////////////////////////////////////////

    let thread_count = num_cpus::get();
    let mut pool = Pool::new(thread_count as u32);

    let capitalSteadyState: f64 = (ALPHA * BETA).powf(1_f64 / (1_f64 - ALPHA));
    let outputSteadyState: f64 = capitalSteadyState.powf(ALPHA);
    let consumptionSteadyState: f64 = outputSteadyState - capitalSteadyState;

    if print {
        println!("Output = {}, Capital = {}, Consumption = {}",
                 outputSteadyState,
                 capitalSteadyState,
                 consumptionSteadyState);
    }

    let mut vGridCapital = [0f64; GRID_CAPITAL];

    for (i, val) in vGridCapital.iter_mut().enumerate() {
        *val = 0.5 * capitalSteadyState + 0.00001 * (i as f64)
    }

    // 3. Required matrices and vectors
    let mut mValueFunction: [Vec<f64>; GRID_PRODUCTIVITY] = [
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
    ];

    let mut mPolicyFunction: [Vec<f64>; GRID_PRODUCTIVITY] = [
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
    ];

    let mut expectedValueFunction: [Vec<f64>; GRID_PRODUCTIVITY] = [
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
    ];

    // 4. We pre-build output for each point in the grid
    let mOutput : Vec<[f64; GRID_PRODUCTIVITY]> = (0..GRID_CAPITAL).map(|nCapital| {
        let mut arr = [0.0; GRID_PRODUCTIVITY];

        for (nProductivity, slot) in arr.iter_mut().enumerate() {
            *slot = PRODUCTIVITY[nProductivity] * vGridCapital[nCapital].powf(ALPHA)
        }

        arr
    }).collect();

  // 5. Main iteration
  // TODO: one could implement a macro for the multiple declarations
    let mut maxDifference = 10_f64;
    const TOLERANCE: f64 = 0.0000001; // compiler warn: variable does not need to be mutable

    let mut iteration = 0;

    // small array to split where the writes are going, so we don't need a mutex
    let mut differences = [-100000.0; GRID_PRODUCTIVITY];

    while maxDifference > TOLERANCE {
        pool.scoped(|scoped| {
            for (expected_values, transitions) in expectedValueFunction.iter_mut()
                                                                       .zip(TRANSITIONS.iter()) {
                // Only capture refs
                let mValueFunction = &mValueFunction;

                scoped.execute(move || {
                    for (idx, expected_value) in expected_values.iter_mut().enumerate() {
                        *expected_value = transitions.iter()
                                                     .zip(mValueFunction.iter())
                                                     .fold(0.0f64, |acc, (transition, value_fns)| {
                            acc + (transition * value_fns[idx])
                        });
                    }
                });
            }
        });

        pool.scoped(|scoped| {
            for (nProductivity, (((policies, value_fns), maxDifference), expected_values)) in mPolicyFunction.iter_mut()
                                                                         .zip(mValueFunction.iter_mut())
                                                                         .zip(differences.iter_mut())
                                                                         .zip(expectedValueFunction.iter())
                                                                         .enumerate() {
                *maxDifference = -100000.0;

                // Only capture refs
                let mOutput = &mOutput;

                scoped.execute(move || {
                    // We start from previous choice (monotonicity of policy function)
                    let mut gridCapitalNextPeriod = 0;

                    for ((policy, output), value_fn) in policies.iter_mut()
                                                                .zip(mOutput.iter())
                                                                .zip(value_fns.iter_mut()) {
                        let mut valueHighSoFar = -100000.0;
                        let mut capitalChoice = vGridCapital[0];
                        let mOutput_cache = &output[nProductivity];

                        for nCapitalNextPeriod in gridCapitalNextPeriod..GRID_CAPITAL {
                            let consumption = mOutput_cache - &vGridCapital[nCapitalNextPeriod];
                            let valueProvisional = (1_f64 - BETA) * (consumption.ln()) +
                                                    BETA *
                                                    expected_values[nCapitalNextPeriod];

                            if valueProvisional > valueHighSoFar {
                                valueHighSoFar = valueProvisional;
                                capitalChoice = vGridCapital[nCapitalNextPeriod];
                                gridCapitalNextPeriod = nCapitalNextPeriod;
                            } else {
                                break; // We break when we have achieved the max
                            }
                        }

                        let old = mem::replace(value_fn, valueHighSoFar);
                        let diff = (old - valueHighSoFar).abs();
                        if diff > *maxDifference {
                            *maxDifference = diff
                        }

                        *policy = capitalChoice;
                    }
                })
            }
        });

        maxDifference = -100000.0;
        for &diff in &differences {
            if diff > maxDifference {
                maxDifference = diff
            }
        }

        iteration += 1;
        if print && (iteration % 10 == 0 || iteration == 1) {
            println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
        }
    }

    if print {
        println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
    }

    mPolicyFunction[2][999]
}

fn main() {
    use std::env::args;

    let sample = args().any(|arg| arg == "--sample");

    if sample {
        let mut samples = (0..5).map(|i| {
            let cpu0 = precise_time_s();
            let result = solve(false);
            let cpu1 = precise_time_s();

            assert_eq!(result, 0.1465491436956954);

            let diff = cpu1 - cpu0;
            println!("Sample #{}, Time: {}s", i + 1, diff);
            diff
        }).collect::<Vec<f64>>();

        samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
        println!("\nMedian time is = {}", samples[2]);
    } else {
        let cpu0 = precise_time_s();
        let result = solve(true);
        let cpu1 = precise_time_s();

        println!("My check = {}\n", result);
        assert_eq!(result, 0.1465491436956954);
        println!("Elapsed time is = {}", cpu1 - cpu0);
    }
}

// FIXME: Can't use bench as the loop is too slow (takes 8mins to get samples)
// #[cfg(test)]
// mod test {
//     extern crate test;
//     use self::test::Bencher;
//     use super::solve;

//     #[bench]
//     fn bench(b: &mut Bencher) {
//         b.iter(|| {
//             // use `test::black_box` to prevent compiler optimizations from disregarding
//             // unused values
//             test::black_box(solve(false));
//         });
//     }
// }
<commit_msg>Rename variables to rust style<commit_after>// #![cfg_attr(test, feature(test))]
extern crate time;
extern crate num_cpus;
extern crate scoped_threadpool;

use time::precise_time_s;
use scoped_threadpool::Pool;
use std::mem;

///////////////////////////////////////////////////////////////////////////////
// 1. Calibration
///////////////////////////////////////////////////////////////////////////////
const ALPHA: f64 = 0.33333333333;  // Elasticity of output w.r.t. capital
const BETA: f64 = 0.95;           // Discount factor

// Productivity values
const PRODUCTIVITY: [f64; 5] = [0.9792, 0.9896, 1.0000, 1.0106, 1.0212];

// Transition matrix
static TRANSITIONS: [[f64; 5]; 5] = [[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],
                                     [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],
                                     [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],
                                     [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],
                                     [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]];

// Dimensions to generate the grid of capital
const GRID_CAPITAL: usize = 17820;
const GRID_PRODUCTIVITY: usize = 5;

fn solve(print: bool) -> f64 {
  /////////////////////////////////////////////////////////////////////////////
  // 2. Steady State
  /////////////////////////////////////////////////////////////////////////////

    let thread_count = num_cpus::get();
    let mut pool = Pool::new(thread_count as u32);

    let capital_steady_state: f64 = (ALPHA * BETA).powf(1_f64 / (1_f64 - ALPHA));
    let output_steady_state: f64 = capital_steady_state.powf(ALPHA);
    let consumption_steady_state: f64 = output_steady_state - capital_steady_state;

    if print {
        println!("Output = {}, Capital = {}, Consumption = {}",
                 output_steady_state,
                 capital_steady_state,
                 consumption_steady_state);
    }

    let mut grid_capital = [0f64; GRID_CAPITAL];

    for (i, val) in grid_capital.iter_mut().enumerate() {
        *val = 0.5 * capital_steady_state + 0.00001 * (i as f64)
    }

    // 3. Required matrices and vectors
    let mut values: [Vec<f64>; GRID_PRODUCTIVITY] = [
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
    ];

    let mut policies: [Vec<f64>; GRID_PRODUCTIVITY] = [
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
    ];

    let mut expected_values: [Vec<f64>; GRID_PRODUCTIVITY] = [
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
        vec![0f64; GRID_CAPITAL],
    ];

    // 4. We pre-build output for each point in the grid
    let output : Vec<[f64; GRID_PRODUCTIVITY]> = (0..GRID_CAPITAL).map(|capital_idx| {
        let mut arr = [0.0; GRID_PRODUCTIVITY];

        for (productivity_idx, slot) in arr.iter_mut().enumerate() {
            *slot = PRODUCTIVITY[productivity_idx] * grid_capital[capital_idx].powf(ALPHA)
        }

        arr
    }).collect();

  // 5. Main iteration
  // TODO: one could implement a macro for the multiple declarations
    let mut max_difference = 10_f64;
    const TOLERANCE: f64 = 0.0000001; // compiler warn: variable does not need to be mutable

    let mut iteration = 0;

    // small array to split where the writes are going, so we don't need a mutex
    let mut diffs = [-100000.0; GRID_PRODUCTIVITY];

    while max_difference > TOLERANCE {
        pool.scoped(|scoped| {
            for (expected_values, transitions) in expected_values.iter_mut()
                                                                       .zip(TRANSITIONS.iter()) {
                // Only capture refs
                let values = &values;

                scoped.execute(move || {
                    for (idx, expected_value) in expected_values.iter_mut().enumerate() {
                        *expected_value = transitions.iter()
                                                     .zip(values.iter())
                                                     .fold(0.0f64, |acc, (transition, value_fns)| {
                            acc + (transition * value_fns[idx])
                        });
                    }
                });
            }
        });

        pool.scoped(|scoped| {
            for (productivity_idx, (((policies, value_fns), max_difference), expected_values)) in policies.iter_mut()
                                                                         .zip(values.iter_mut())
                                                                         .zip(diffs.iter_mut())
                                                                         .zip(expected_values.iter())
                                                                         .enumerate() {
                *max_difference = -100000.0;

                // Only capture refs
                let output = &output;

                scoped.execute(move || {
                    // We start from previous choice (monotonicity of policy function)
                    let mut grid_capital_next_period = 0;

                    for ((policy, output), value_fn) in policies.iter_mut()
                                                                .zip(output.iter())
                                                                .zip(value_fns.iter_mut()) {
                        let mut value_high = -100000.0;
                        let mut capital_choice = grid_capital[0];
                        let output_cache = &output[productivity_idx];

                        for capital_next_period in grid_capital_next_period..GRID_CAPITAL {
                            let consumption = output_cache - &grid_capital[capital_next_period];
                            let value_provisional = (1_f64 - BETA) * (consumption.ln()) +
                                                    BETA *
                                                    expected_values[capital_next_period];

                            if value_provisional > value_high {
                                value_high = value_provisional;
                                capital_choice = grid_capital[capital_next_period];
                                grid_capital_next_period = capital_next_period;
                            } else {
                                break; // We break when we have achieved the max
                            }
                        }

                        let old = mem::replace(value_fn, value_high);
                        let diff = (old - value_high).abs();
                        if diff > *max_difference {
                            *max_difference = diff
                        }

                        *policy = capital_choice;
                    }
                })
            }
        });

        max_difference = -100000.0;
        for &diff in &diffs {
            if diff > max_difference {
                max_difference = diff
            }
        }

        iteration += 1;
        if print && (iteration % 10 == 0 || iteration == 1) {
            println!("Iteration = {}, Sup Diff = {}", iteration, max_difference);
        }
    }

    if print {
        println!("Iteration = {}, Sup Diff = {}", iteration, max_difference);
    }

    policies[2][999]
}

fn main() {
    use std::env::args;

    let sample = args().any(|arg| arg == "--sample");

    if sample {
        let mut samples = (0..5).map(|i| {
            let cpu0 = precise_time_s();
            let result = solve(false);
            let cpu1 = precise_time_s();

            assert_eq!(result, 0.1465491436956954);

            let diff = cpu1 - cpu0;
            println!("Sample #{}, Time: {}s", i + 1, diff);
            diff
        }).collect::<Vec<f64>>();

        samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
        println!("\nMedian time is = {}", samples[2]);
    } else {
        let cpu0 = precise_time_s();
        let result = solve(true);
        let cpu1 = precise_time_s();

        println!("My check = {}\n", result);
        assert_eq!(result, 0.1465491436956954);
        println!("Elapsed time is = {}", cpu1 - cpu0);
    }
}

// FIXME: Can't use bench as the loop is too slow (takes 8mins to get samples)
// #[cfg(test)]
// mod test {
//     extern crate test;
//     use self::test::Bencher;
//     use super::solve;

//     #[bench]
//     fn bench(b: &mut Bencher) {
//         b.iter(|| {
//             // use `test::black_box` to prevent compiler optimizations from disregarding
//             // unused values
//             test::black_box(solve(false));
//         });
//     }
// }
<|endoftext|>