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

//! This crate provides the Faucet service for creating and funding accounts on the Aptos Network.
//!
//! THIS SERVICE SHOULD NEVER BE DEPLOYED TO MAINNET.
//!
//! ## Launch service
//!
//! Launch faucet service locally and connect to Testnet:
//!
//! ```bash
//! cargo run --bin aptos-faucet -- -c TESTNET -m <mint-private-key-path> -s http://localhost:8080 -p 8081
//! ```
//!
//! Check help doc for options details:
//!
//! ```bash
//! cargo run -p aptos-faucet -- -h
//! ```

use anyhow::{anyhow, Result};
use aptos_crypto::ed25519::Ed25519PublicKey;
use aptos_logger::info;
use aptos_rest_client::Client;
use aptos_sdk::{
    transaction_builder::{aptos_stdlib, TransactionFactory},
    types::{
        account_address::AccountAddress,
        chain_id::ChainId,
        transaction::{
            authenticator::{AuthenticationKey, AuthenticationKeyPreimage},
            SignedTransaction,
        },
        LocalAccount,
    },
};
use reqwest::StatusCode;
use serde::Deserialize;
use std::{
    convert::Infallible,
    fmt,
    sync::{Arc, Mutex},
};
use url::Url;
use warp::{Filter, Rejection, Reply};

pub mod mint;

pub struct Service {
    pub faucet_account: Mutex<LocalAccount>,
    transaction_factory: TransactionFactory,
    client: Client,
    endpoint: String,
    maximum_amount: Option<u64>,
}

impl Service {
    pub fn new(
        endpoint: String,
        chain_id: ChainId,
        faucet_account: LocalAccount,
        maximum_amount: Option<u64>,
    ) -> Self {
        let client = Client::new(Url::parse(&endpoint).expect("Invalid rest endpoint"));
        Service {
            faucet_account: Mutex::new(faucet_account),
            transaction_factory: TransactionFactory::new(chain_id)
                .with_transaction_expiration_time(30),
            client,
            endpoint,
            maximum_amount,
        }
    }

    pub fn endpoint(&self) -> &String {
        &self.endpoint
    }
}

pub fn routes(
    service: Arc<Service>,
) -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone {
    let mint = mint::mint_routes(service.clone());
    let accounts = accounts_routes(service.clone());
    let health = health_route(service);

    health
        .or(mint)
        .or(accounts)
        .with(warp::log::custom(|info| {
            info!(
                "{} \"{} {} {:?}\" {} \"{}\" \"{}\" {:?}",
                OptFmt(info.remote_addr()),
                info.method(),
                info.path(),
                info.version(),
                info.status().as_u16(),
                OptFmt(info.referer()),
                OptFmt(info.user_agent()),
                info.elapsed(),
            )
        }))
        .with(warp::cors().allow_any_origin().allow_methods(vec!["POST"]))
}

fn health_route(
    service: Arc<Service>,
) -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone {
    warp::path!("health")
        .and(warp::get())
        .and(warp::any().map(move || service.clone()))
        .and_then(handle_health)
}

async fn handle_health(service: Arc<Service>) -> Result<Box<dyn warp::Reply>, Infallible> {
    let faucet_address = service.faucet_account.lock().unwrap().address();
    let faucet_account = service.client.get_account(faucet_address).await;

    match faucet_account {
        Ok(account) => Ok(Box::new(account.inner().sequence_number.to_string())),
        Err(err) => Ok(Box::new(warp::reply::with_status(
            err.to_string(),
            StatusCode::INTERNAL_SERVER_ERROR,
        ))),
    }
}

fn accounts_routes(
    service: Arc<Service>,
) -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone {
    create_account_route(service.clone()).or(fund_account_route(service))
}

#[derive(Deserialize)]
struct CreateAccountParams {
    pub_key: Ed25519PublicKey,
}

impl CreateAccountParams {
    fn pre_image(&self) -> AuthenticationKeyPreimage {
        AuthenticationKeyPreimage::ed25519(&self.pub_key)
    }

    fn receiver(&self) -> AccountAddress {
        AuthenticationKey::ed25519(&self.pub_key).derived_address()
    }
}
fn create_account_route(
    service: Arc<Service>,
) -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone {
    warp::path::path("accounts")
        .and(warp::post())
        .and(warp::any().map(move || service.clone()))
        .and(warp::query().map(move |params: CreateAccountParams| params))
        .and_then(handle_create_account)
}

async fn handle_create_account(
    service: Arc<Service>,
    params: CreateAccountParams,
) -> Result<Box<dyn warp::Reply>, Infallible> {
    match create_account(service, params).await {
        Ok(txn) => Ok(Box::new(hex::encode(bcs::to_bytes(&txn).unwrap()))),
        Err(err) => Ok(Box::new(warp::reply::with_status(
            err.to_string(),
            StatusCode::INTERNAL_SERVER_ERROR,
        ))),
    }
}

async fn create_account(
    service: Arc<Service>,
    params: CreateAccountParams,
) -> Result<SignedTransaction> {
    // Check to ensure the account hasn't already been created
    if service.client.get_account(params.receiver()).await.is_ok() {
        return Err(anyhow!("account already exists"));
    }

    let faucet_account_address = service.faucet_account.lock().unwrap().address();
    let faucet_sequence_number = service
        .client
        .get_account(faucet_account_address)
        .await
        .map_err(|_| anyhow::format_err!("faucet account {} not found", faucet_account_address))?
        .into_inner()
        .sequence_number;

    let txn = {
        let mut faucet_account = service.faucet_account.lock().unwrap();
        if faucet_sequence_number > faucet_account.sequence_number() {
            *faucet_account.sequence_number_mut() = faucet_sequence_number;
        }

        let builder = service.transaction_factory.payload(
            aptos_stdlib::encode_create_account_script_function(
                params.receiver(),
                params.pre_image().into_vec(),
            ),
        );

        faucet_account.sign_with_transaction_builder(builder)
    };

    service.client.submit(&txn).await?;
    Ok(txn)
}

#[derive(Deserialize)]
struct FundAccountParams {
    amount: Option<u64>,
}

fn fund_account_route(
    service: Arc<Service>,
) -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone {
    warp::path!("accounts" / AccountAddress / "fund")
        .and(warp::post())
        .and(warp::any().map(move || service.clone()))
        .and(warp::query().map(move |params: FundAccountParams| params))
        .and_then(handle_fund_account)
}

async fn handle_fund_account(
    address: AccountAddress,
    service: Arc<Service>,
    params: FundAccountParams,
) -> Result<Box<dyn warp::Reply>, Infallible> {
    match fund_account(service, address, params).await {
        Ok(txn) => Ok(Box::new(hex::encode(bcs::to_bytes(&txn).unwrap()))),
        Err(err) => Ok(Box::new(warp::reply::with_status(
            err.to_string(),
            StatusCode::INTERNAL_SERVER_ERROR,
        ))),
    }
}

async fn fund_account(
    service: Arc<Service>,
    address: AccountAddress,
    params: FundAccountParams,
) -> Result<SignedTransaction> {
    let asked_amount = params
        .amount
        .ok_or_else(|| anyhow::format_err!("Mint amount must be provided"))?;
    let service_amount = service.maximum_amount.unwrap_or(asked_amount);
    let amount = std::cmp::min(asked_amount, service_amount);

    // Check to ensure the account has already been created
    if service.client.get_account(address).await.is_err() {
        return Err(anyhow!("account doesn't exist"));
    }

    let faucet_account_address = service.faucet_account.lock().unwrap().address();
    let faucet_sequence_number = service
        .client
        .get_account(faucet_account_address)
        .await
        .map_err(|_| anyhow::format_err!("faucet account {} not found", faucet_account_address))?
        .into_inner()
        .sequence_number;

    let txn = {
        let mut faucet_account = service.faucet_account.lock().unwrap();
        if faucet_sequence_number > faucet_account.sequence_number() {
            *faucet_account.sequence_number_mut() = faucet_sequence_number;
        }

        faucet_account.sign_with_transaction_builder(
            service
                .transaction_factory
                .payload(aptos_stdlib::encode_mint_script_function(address, amount)),
        )
    };

    service.client.submit(&txn).await?;
    Ok(txn)
}

//
// Common Types
//

struct OptFmt<T>(Option<T>);

impl<T: fmt::Display> fmt::Display for OptFmt<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Some(t) = &self.0 {
            fmt::Display::fmt(t, f)
        } else {
            f.write_str("-")
        }
    }
}
<commit_msg>[faucet] specify a gas base unit for testnet<commit_after>// Copyright (c) The Aptos Foundation
// SPDX-License-Identifier: Apache-2.0

//! This crate provides the Faucet service for creating and funding accounts on the Aptos Network.
//!
//! THIS SERVICE SHOULD NEVER BE DEPLOYED TO MAINNET.
//!
//! ## Launch service
//!
//! Launch faucet service locally and connect to Testnet:
//!
//! ```bash
//! cargo run --bin aptos-faucet -- -c TESTNET -m <mint-private-key-path> -s http://localhost:8080 -p 8081
//! ```
//!
//! Check help doc for options details:
//!
//! ```bash
//! cargo run -p aptos-faucet -- -h
//! ```

use anyhow::{anyhow, Result};
use aptos_crypto::ed25519::Ed25519PublicKey;
use aptos_logger::info;
use aptos_rest_client::Client;
use aptos_sdk::{
    transaction_builder::{aptos_stdlib, TransactionFactory},
    types::{
        account_address::AccountAddress,
        chain_id::ChainId,
        transaction::{
            authenticator::{AuthenticationKey, AuthenticationKeyPreimage},
            SignedTransaction,
        },
        LocalAccount,
    },
};
use reqwest::StatusCode;
use serde::Deserialize;
use std::{
    convert::Infallible,
    fmt,
    sync::{Arc, Mutex},
};
use url::Url;
use warp::{Filter, Rejection, Reply};

pub mod mint;

pub struct Service {
    pub faucet_account: Mutex<LocalAccount>,
    transaction_factory: TransactionFactory,
    client: Client,
    endpoint: String,
    maximum_amount: Option<u64>,
}

impl Service {
    pub fn new(
        endpoint: String,
        chain_id: ChainId,
        faucet_account: LocalAccount,
        maximum_amount: Option<u64>,
    ) -> Self {
        let client = Client::new(Url::parse(&endpoint).expect("Invalid rest endpoint"));
        Service {
            faucet_account: Mutex::new(faucet_account),
            transaction_factory: TransactionFactory::new(chain_id)
                .with_gas_unit_price(1)
                .with_transaction_expiration_time(30),
            client,
            endpoint,
            maximum_amount,
        }
    }

    pub fn endpoint(&self) -> &String {
        &self.endpoint
    }
}

pub fn routes(
    service: Arc<Service>,
) -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone {
    let mint = mint::mint_routes(service.clone());
    let accounts = accounts_routes(service.clone());
    let health = health_route(service);

    health
        .or(mint)
        .or(accounts)
        .with(warp::log::custom(|info| {
            info!(
                "{} \"{} {} {:?}\" {} \"{}\" \"{}\" {:?}",
                OptFmt(info.remote_addr()),
                info.method(),
                info.path(),
                info.version(),
                info.status().as_u16(),
                OptFmt(info.referer()),
                OptFmt(info.user_agent()),
                info.elapsed(),
            )
        }))
        .with(warp::cors().allow_any_origin().allow_methods(vec!["POST"]))
}

fn health_route(
    service: Arc<Service>,
) -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone {
    warp::path!("health")
        .and(warp::get())
        .and(warp::any().map(move || service.clone()))
        .and_then(handle_health)
}

async fn handle_health(service: Arc<Service>) -> Result<Box<dyn warp::Reply>, Infallible> {
    let faucet_address = service.faucet_account.lock().unwrap().address();
    let faucet_account = service.client.get_account(faucet_address).await;

    match faucet_account {
        Ok(account) => Ok(Box::new(account.inner().sequence_number.to_string())),
        Err(err) => Ok(Box::new(warp::reply::with_status(
            err.to_string(),
            StatusCode::INTERNAL_SERVER_ERROR,
        ))),
    }
}

fn accounts_routes(
    service: Arc<Service>,
) -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone {
    create_account_route(service.clone()).or(fund_account_route(service))
}

#[derive(Deserialize)]
struct CreateAccountParams {
    pub_key: Ed25519PublicKey,
}

impl CreateAccountParams {
    fn pre_image(&self) -> AuthenticationKeyPreimage {
        AuthenticationKeyPreimage::ed25519(&self.pub_key)
    }

    fn receiver(&self) -> AccountAddress {
        AuthenticationKey::ed25519(&self.pub_key).derived_address()
    }
}
fn create_account_route(
    service: Arc<Service>,
) -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone {
    warp::path::path("accounts")
        .and(warp::post())
        .and(warp::any().map(move || service.clone()))
        .and(warp::query().map(move |params: CreateAccountParams| params))
        .and_then(handle_create_account)
}

async fn handle_create_account(
    service: Arc<Service>,
    params: CreateAccountParams,
) -> Result<Box<dyn warp::Reply>, Infallible> {
    match create_account(service, params).await {
        Ok(txn) => Ok(Box::new(hex::encode(bcs::to_bytes(&txn).unwrap()))),
        Err(err) => Ok(Box::new(warp::reply::with_status(
            err.to_string(),
            StatusCode::INTERNAL_SERVER_ERROR,
        ))),
    }
}

async fn create_account(
    service: Arc<Service>,
    params: CreateAccountParams,
) -> Result<SignedTransaction> {
    // Check to ensure the account hasn't already been created
    if service.client.get_account(params.receiver()).await.is_ok() {
        return Err(anyhow!("account already exists"));
    }

    let faucet_account_address = service.faucet_account.lock().unwrap().address();
    let faucet_sequence_number = service
        .client
        .get_account(faucet_account_address)
        .await
        .map_err(|_| anyhow::format_err!("faucet account {} not found", faucet_account_address))?
        .into_inner()
        .sequence_number;

    let txn = {
        let mut faucet_account = service.faucet_account.lock().unwrap();
        if faucet_sequence_number > faucet_account.sequence_number() {
            *faucet_account.sequence_number_mut() = faucet_sequence_number;
        }

        let builder = service.transaction_factory.payload(
            aptos_stdlib::encode_create_account_script_function(
                params.receiver(),
                params.pre_image().into_vec(),
            ),
        );

        faucet_account.sign_with_transaction_builder(builder)
    };

    service.client.submit(&txn).await?;
    Ok(txn)
}

#[derive(Deserialize)]
struct FundAccountParams {
    amount: Option<u64>,
}

fn fund_account_route(
    service: Arc<Service>,
) -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone {
    warp::path!("accounts" / AccountAddress / "fund")
        .and(warp::post())
        .and(warp::any().map(move || service.clone()))
        .and(warp::query().map(move |params: FundAccountParams| params))
        .and_then(handle_fund_account)
}

async fn handle_fund_account(
    address: AccountAddress,
    service: Arc<Service>,
    params: FundAccountParams,
) -> Result<Box<dyn warp::Reply>, Infallible> {
    match fund_account(service, address, params).await {
        Ok(txn) => Ok(Box::new(hex::encode(bcs::to_bytes(&txn).unwrap()))),
        Err(err) => Ok(Box::new(warp::reply::with_status(
            err.to_string(),
            StatusCode::INTERNAL_SERVER_ERROR,
        ))),
    }
}

async fn fund_account(
    service: Arc<Service>,
    address: AccountAddress,
    params: FundAccountParams,
) -> Result<SignedTransaction> {
    let asked_amount = params
        .amount
        .ok_or_else(|| anyhow::format_err!("Mint amount must be provided"))?;
    let service_amount = service.maximum_amount.unwrap_or(asked_amount);
    let amount = std::cmp::min(asked_amount, service_amount);

    // Check to ensure the account has already been created
    if service.client.get_account(address).await.is_err() {
        return Err(anyhow!("account doesn't exist"));
    }

    let faucet_account_address = service.faucet_account.lock().unwrap().address();
    let faucet_sequence_number = service
        .client
        .get_account(faucet_account_address)
        .await
        .map_err(|_| anyhow::format_err!("faucet account {} not found", faucet_account_address))?
        .into_inner()
        .sequence_number;

    let txn = {
        let mut faucet_account = service.faucet_account.lock().unwrap();
        if faucet_sequence_number > faucet_account.sequence_number() {
            *faucet_account.sequence_number_mut() = faucet_sequence_number;
        }

        faucet_account.sign_with_transaction_builder(
            service
                .transaction_factory
                .payload(aptos_stdlib::encode_mint_script_function(address, amount)),
        )
    };

    service.client.submit(&txn).await?;
    Ok(txn)
}

//
// Common Types
//

struct OptFmt<T>(Option<T>);

impl<T: fmt::Display> fmt::Display for OptFmt<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Some(t) = &self.0 {
            fmt::Display::fmt(t, f)
        } else {
            f.write_str("-")
        }
    }
}
<|endoftext|>


<commit_before>use hir::Semantics;
use ide_db::{
    base_db::FileId, helpers::pick_best_token,
    syntax_helpers::insert_whitespace_into_node::insert_ws_into, RootDatabase,
};
use syntax::{ast, ted, AstNode, NodeOrToken, SyntaxKind, SyntaxNode, T};

use crate::FilePosition;

pub struct ExpandedMacro {
    pub name: String,
    pub expansion: String,
}

// Feature: Expand Macro Recursively
//
// Shows the full macro expansion of the macro at current cursor.
//
// |===
// | Editor  | Action Name
//
// | VS Code | **Rust Analyzer: Expand macro recursively**
// |===
//
// image::https://user-images.githubusercontent.com/48062697/113020648-b3973180-917a-11eb-84a9-ecb921293dc5.gif[]
pub(crate) fn expand_macro(db: &RootDatabase, position: FilePosition) -> Option<ExpandedMacro> {
    let sema = Semantics::new(db);
    let file = sema.parse(position.file_id);

    let tok = pick_best_token(file.syntax().token_at_offset(position.offset), |kind| match kind {
        SyntaxKind::IDENT => 1,
        _ => 0,
    })?;

    // due to how Rust Analyzer works internally, we need to special case derive attributes,
    // otherwise they might not get found, e.g. here with the cursor at $0 `#[attr]` would expand:
    // ```
    // #[attr]
    // #[derive($0Foo)]
    // struct Bar;
    // ```

    let derive = sema.descend_into_macros(tok.clone()).into_iter().find_map(|descended| {
        let hir_file = sema.hir_file_for(&descended.parent()?);
        if !hir_file.is_derive_attr_pseudo_expansion(db) {
            return None;
        }

        let name = descended.ancestors().filter_map(ast::Path::cast).last()?.to_string();
        // up map out of the #[derive] expansion
        let token = hir::InFile::new(hir_file, descended).upmap(db)?.value;
        let attr = token.ancestors().find_map(ast::Attr::cast)?;
        let expansions = sema.expand_derive_macro(&attr)?;
        let idx = attr
            .token_tree()?
            .token_trees_and_tokens()
            .filter_map(NodeOrToken::into_token)
            .take_while(|it| it != &token)
            .filter(|it| it.kind() == T![,])
            .count();
        let expansion =
            format(db, SyntaxKind::MACRO_ITEMS, position.file_id, expansions.get(idx).cloned()?);
        Some(ExpandedMacro { name, expansion })
    });

    if derive.is_some() {
        return derive;
    }

    // FIXME: Intermix attribute and bang! expansions
    // currently we only recursively expand one of the two types
    let mut anc = tok.ancestors();
    let (name, expanded, kind) = loop {
        let node = anc.next()?;

        if let Some(item) = ast::Item::cast(node.clone()) {
            if let Some(def) = sema.resolve_attr_macro_call(&item) {
                break (
                    def.name(db).to_string(),
                    expand_attr_macro_recur(&sema, &item)?,
                    SyntaxKind::MACRO_ITEMS,
                );
            }
        }
        if let Some(mac) = ast::MacroCall::cast(node) {
            break (
                mac.path()?.segment()?.name_ref()?.to_string(),
                expand_macro_recur(&sema, &mac)?,
                mac.syntax().parent().map(|it| it.kind()).unwrap_or(SyntaxKind::MACRO_ITEMS),
            );
        }
    };

    // FIXME:
    // macro expansion may lose all white space information
    // But we hope someday we can use ra_fmt for that
    let expansion = format(db, kind, position.file_id, expanded);

    Some(ExpandedMacro { name, expansion })
}

fn expand_macro_recur(
    sema: &Semantics<RootDatabase>,
    macro_call: &ast::MacroCall,
) -> Option<SyntaxNode> {
    let expanded = sema.expand(macro_call)?.clone_for_update();
    expand(sema, expanded, ast::MacroCall::cast, expand_macro_recur)
}

fn expand_attr_macro_recur(sema: &Semantics<RootDatabase>, item: &ast::Item) -> Option<SyntaxNode> {
    let expanded = sema.expand_attr_macro(item)?.clone_for_update();
    expand(sema, expanded, ast::Item::cast, expand_attr_macro_recur)
}

fn expand<T: AstNode>(
    sema: &Semantics<RootDatabase>,
    expanded: SyntaxNode,
    f: impl FnMut(SyntaxNode) -> Option<T>,
    exp: impl Fn(&Semantics<RootDatabase>, &T) -> Option<SyntaxNode>,
) -> Option<SyntaxNode> {
    let children = expanded.descendants().filter_map(f);
    let mut replacements = Vec::new();

    for child in children {
        if let Some(new_node) = exp(sema, &child) {
            // check if the whole original syntax is replaced
            if expanded == *child.syntax() {
                return Some(new_node);
            }
            replacements.push((child, new_node));
        }
    }

    replacements.into_iter().rev().for_each(|(old, new)| ted::replace(old.syntax(), new));
    Some(expanded)
}

fn format(db: &RootDatabase, kind: SyntaxKind, file_id: FileId, expanded: SyntaxNode) -> String {
    let expansion = insert_ws_into(expanded).to_string();

    _format(db, kind, file_id, &expansion).unwrap_or(expansion)
}

#[cfg(any(test, target_arch = "wasm32", target_os = "emscripten"))]
fn _format(
    _db: &RootDatabase,
    _kind: SyntaxKind,
    _file_id: FileId,
    _expansion: &str,
) -> Option<String> {
    None
}

#[cfg(not(any(test, target_arch = "wasm32", target_os = "emscripten")))]
fn _format(
    db: &RootDatabase,
    kind: SyntaxKind,
    file_id: FileId,
    expansion: &str,
) -> Option<String> {
    use ide_db::base_db::{FileLoader, SourceDatabase};
    // hack until we get hygiene working (same character amount to preserve formatting as much as possible)
    const DOLLAR_CRATE_REPLACE: &str = &"__r_a_";
    let expansion = expansion.replace("$crate", DOLLAR_CRATE_REPLACE);
    let (prefix, suffix) = match kind {
        SyntaxKind::MACRO_PAT => ("fn __(", ": u32);"),
        SyntaxKind::MACRO_EXPR | SyntaxKind::MACRO_STMTS => ("fn __() {", "}"),
        SyntaxKind::MACRO_TYPE => ("type __ =", ";"),
        _ => ("", ""),
    };
    let expansion = format!("{prefix}{expansion}{suffix}");

    let &crate_id = db.relevant_crates(file_id).iter().next()?;
    let edition = db.crate_graph()[crate_id].edition;

    let mut cmd = std::process::Command::new(toolchain::rustfmt());
    cmd.arg("--edition");
    cmd.arg(edition.to_string());

    let mut rustfmt = cmd
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .spawn()
        .ok()?;

    std::io::Write::write_all(&mut rustfmt.stdin.as_mut()?, expansion.as_bytes()).ok()?;

    let output = rustfmt.wait_with_output().ok()?;
    let captured_stdout = String::from_utf8(output.stdout).ok()?;

    if output.status.success() && !captured_stdout.trim().is_empty() {
        let output = captured_stdout.replace(DOLLAR_CRATE_REPLACE, "$crate");
        let output = output.trim().strip_prefix(prefix)?;
        let output = match kind {
            SyntaxKind::MACRO_PAT => {
                output.strip_suffix(suffix).or_else(|| output.strip_suffix(": u32,\n);"))?
            }
            _ => output.strip_suffix(suffix)?,
        };
        let trim_indent = stdx::trim_indent(output);
        tracing::debug!("expand_macro: formatting succeeded");
        Some(trim_indent)
    } else {
        None
    }
}

#[cfg(test)]
mod tests {
    use expect_test::{expect, Expect};

    use crate::fixture;

    #[track_caller]
    fn check(ra_fixture: &str, expect: Expect) {
        let (analysis, pos) = fixture::position(ra_fixture);
        let expansion = analysis.expand_macro(pos).unwrap().unwrap();
        let actual = format!("{}\n{}", expansion.name, expansion.expansion);
        expect.assert_eq(&actual);
    }

    #[test]
    fn macro_expand_as_keyword() {
        check(
            r#"
macro_rules! bar {
    ($i:tt) => { $i as _ }
}
fn main() {
    let x: u64 = ba$0r!(5i64);
}
"#,
            expect![[r#"
                bar
                5i64 as _"#]],
        );
    }

    #[test]
    fn macro_expand_recursive_expansion() {
        check(
            r#"
macro_rules! bar {
    () => { fn  b() {} }
}
macro_rules! foo {
    () => { bar!(); }
}
macro_rules! baz {
    () => { foo!(); }
}
f$0oo!();
"#,
            expect![[r#"
                foo
                fn b(){}

            "#]],
        );
    }

    #[test]
    fn macro_expand_multiple_lines() {
        check(
            r#"
macro_rules! foo {
    () => {
        fn some_thing() -> u32 {
            let a = 0;
            a + 10
        }
    }
}
f$0oo!();
        "#,
            expect![[r#"
                foo
                fn some_thing() -> u32 {
                  let a = 0;
                  a+10
                }
            "#]],
        );
    }

    #[test]
    fn macro_expand_match_ast() {
        check(
            r#"
macro_rules! match_ast {
    (match $node:ident { $($tt:tt)* }) => { match_ast!(match ($node) { $($tt)* }) };
    (match ($node:expr) {
        $( ast::$ast:ident($it:ident) => $res:block, )*
        _ => $catch_all:expr $(,)?
    }) => {{
        $( if let Some($it) = ast::$ast::cast($node.clone()) $res else )*
        { $catch_all }
    }};
}

fn main() {
    mat$0ch_ast! {
        match container {
            ast::TraitDef(it) => {},
            ast::ImplDef(it) => {},
            _ => { continue },
        }
    }
}
"#,
            expect![[r#"
       match_ast
       {
         if let Some(it) = ast::TraitDef::cast(container.clone()){}
         else if let Some(it) = ast::ImplDef::cast(container.clone()){}
         else {
           {
             continue
           }
         }
       }"#]],
        );
    }

    #[test]
    fn macro_expand_match_ast_inside_let_statement() {
        check(
            r#"
macro_rules! match_ast {
    (match $node:ident { $($tt:tt)* }) => { match_ast!(match ($node) { $($tt)* }) };
    (match ($node:expr) {}) => {{}};
}

fn main() {
    let p = f(|it| {
        let res = mat$0ch_ast! { match c {}};
        Some(res)
    })?;
}
"#,
            expect![[r#"
                match_ast
                {}
            "#]],
        );
    }

    #[test]
    fn macro_expand_inner_macro_fail_to_expand() {
        check(
            r#"
macro_rules! bar {
    (BAD) => {};
}
macro_rules! foo {
    () => {bar!()};
}

fn main() {
    let res = fo$0o!();
}
"#,
            expect![[r#"
                foo
            "#]],
        );
    }

    #[test]
    fn macro_expand_with_dollar_crate() {
        check(
            r#"
#[macro_export]
macro_rules! bar {
    () => {0};
}
macro_rules! foo {
    () => {$crate::bar!()};
}

fn main() {
    let res = fo$0o!();
}
"#,
            expect![[r#"
                foo
                0"#]],
        );
    }

    #[test]
    fn macro_expand_with_dyn_absolute_path() {
        check(
            r#"
macro_rules! foo {
    () => {fn f<T>(_: &dyn ::std::marker::Copy) {}};
}

fn main() {
    let res = fo$0o!();
}
"#,
            expect![[r#"
                foo
                fn f<T>(_: &dyn ::std::marker::Copy){}
            "#]],
        );
    }

    #[test]
    fn macro_expand_derive() {
        check(
            r#"
//- proc_macros: identity
//- minicore: clone, derive

#[proc_macros::identity]
#[derive(C$0lone)]
struct Foo {}
"#,
            expect![[r#"
                Clone
                impl < >core::clone::Clone for Foo< >{}

            "#]],
        );
    }

    #[test]
    fn macro_expand_derive2() {
        check(
            r#"
//- minicore: copy, clone, derive

#[derive(Cop$0y)]
#[derive(Clone)]
struct Foo {}
"#,
            expect![[r#"
                Copy
                impl < >core::marker::Copy for Foo< >{}

            "#]],
        );
    }

    #[test]
    fn macro_expand_derive_multi() {
        check(
            r#"
//- minicore: copy, clone, derive

#[derive(Cop$0y, Clone)]
struct Foo {}
"#,
            expect![[r#"
                Copy
                impl < >core::marker::Copy for Foo< >{}

            "#]],
        );
        check(
            r#"
//- minicore: copy, clone, derive

#[derive(Copy, Cl$0one)]
struct Foo {}
"#,
            expect![[r#"
                Clone
                impl < >core::clone::Clone for Foo< >{}

            "#]],
        );
    }
}
<commit_msg>add test for macro expand formatting<commit_after>use hir::Semantics;
use ide_db::{
    base_db::FileId, helpers::pick_best_token,
    syntax_helpers::insert_whitespace_into_node::insert_ws_into, RootDatabase,
};
use syntax::{ast, ted, AstNode, NodeOrToken, SyntaxKind, SyntaxNode, T};

use crate::FilePosition;

pub struct ExpandedMacro {
    pub name: String,
    pub expansion: String,
}

// Feature: Expand Macro Recursively
//
// Shows the full macro expansion of the macro at current cursor.
//
// |===
// | Editor  | Action Name
//
// | VS Code | **Rust Analyzer: Expand macro recursively**
// |===
//
// image::https://user-images.githubusercontent.com/48062697/113020648-b3973180-917a-11eb-84a9-ecb921293dc5.gif[]
pub(crate) fn expand_macro(db: &RootDatabase, position: FilePosition) -> Option<ExpandedMacro> {
    let sema = Semantics::new(db);
    let file = sema.parse(position.file_id);

    let tok = pick_best_token(file.syntax().token_at_offset(position.offset), |kind| match kind {
        SyntaxKind::IDENT => 1,
        _ => 0,
    })?;

    // due to how Rust Analyzer works internally, we need to special case derive attributes,
    // otherwise they might not get found, e.g. here with the cursor at $0 `#[attr]` would expand:
    // ```
    // #[attr]
    // #[derive($0Foo)]
    // struct Bar;
    // ```

    let derive = sema.descend_into_macros(tok.clone()).into_iter().find_map(|descended| {
        let hir_file = sema.hir_file_for(&descended.parent()?);
        if !hir_file.is_derive_attr_pseudo_expansion(db) {
            return None;
        }

        let name = descended.ancestors().filter_map(ast::Path::cast).last()?.to_string();
        // up map out of the #[derive] expansion
        let token = hir::InFile::new(hir_file, descended).upmap(db)?.value;
        let attr = token.ancestors().find_map(ast::Attr::cast)?;
        let expansions = sema.expand_derive_macro(&attr)?;
        let idx = attr
            .token_tree()?
            .token_trees_and_tokens()
            .filter_map(NodeOrToken::into_token)
            .take_while(|it| it != &token)
            .filter(|it| it.kind() == T![,])
            .count();
        let expansion =
            format(db, SyntaxKind::MACRO_ITEMS, position.file_id, expansions.get(idx).cloned()?);
        Some(ExpandedMacro { name, expansion })
    });

    if derive.is_some() {
        return derive;
    }

    // FIXME: Intermix attribute and bang! expansions
    // currently we only recursively expand one of the two types
    let mut anc = tok.ancestors();
    let (name, expanded, kind) = loop {
        let node = anc.next()?;

        if let Some(item) = ast::Item::cast(node.clone()) {
            if let Some(def) = sema.resolve_attr_macro_call(&item) {
                break (
                    def.name(db).to_string(),
                    expand_attr_macro_recur(&sema, &item)?,
                    SyntaxKind::MACRO_ITEMS,
                );
            }
        }
        if let Some(mac) = ast::MacroCall::cast(node) {
            break (
                mac.path()?.segment()?.name_ref()?.to_string(),
                expand_macro_recur(&sema, &mac)?,
                mac.syntax().parent().map(|it| it.kind()).unwrap_or(SyntaxKind::MACRO_ITEMS),
            );
        }
    };

    // FIXME:
    // macro expansion may lose all white space information
    // But we hope someday we can use ra_fmt for that
    let expansion = format(db, kind, position.file_id, expanded);

    Some(ExpandedMacro { name, expansion })
}

fn expand_macro_recur(
    sema: &Semantics<RootDatabase>,
    macro_call: &ast::MacroCall,
) -> Option<SyntaxNode> {
    let expanded = sema.expand(macro_call)?.clone_for_update();
    expand(sema, expanded, ast::MacroCall::cast, expand_macro_recur)
}

fn expand_attr_macro_recur(sema: &Semantics<RootDatabase>, item: &ast::Item) -> Option<SyntaxNode> {
    let expanded = sema.expand_attr_macro(item)?.clone_for_update();
    expand(sema, expanded, ast::Item::cast, expand_attr_macro_recur)
}

fn expand<T: AstNode>(
    sema: &Semantics<RootDatabase>,
    expanded: SyntaxNode,
    f: impl FnMut(SyntaxNode) -> Option<T>,
    exp: impl Fn(&Semantics<RootDatabase>, &T) -> Option<SyntaxNode>,
) -> Option<SyntaxNode> {
    let children = expanded.descendants().filter_map(f);
    let mut replacements = Vec::new();

    for child in children {
        if let Some(new_node) = exp(sema, &child) {
            // check if the whole original syntax is replaced
            if expanded == *child.syntax() {
                return Some(new_node);
            }
            replacements.push((child, new_node));
        }
    }

    replacements.into_iter().rev().for_each(|(old, new)| ted::replace(old.syntax(), new));
    Some(expanded)
}

fn format(db: &RootDatabase, kind: SyntaxKind, file_id: FileId, expanded: SyntaxNode) -> String {
    let expansion = insert_ws_into(expanded).to_string();

    _format(db, kind, file_id, &expansion).unwrap_or(expansion)
}

#[cfg(any(test, target_arch = "wasm32", target_os = "emscripten"))]
fn _format(
    _db: &RootDatabase,
    _kind: SyntaxKind,
    _file_id: FileId,
    _expansion: &str,
) -> Option<String> {
    None
}

#[cfg(not(any(test, target_arch = "wasm32", target_os = "emscripten")))]
fn _format(
    db: &RootDatabase,
    kind: SyntaxKind,
    file_id: FileId,
    expansion: &str,
) -> Option<String> {
    use ide_db::base_db::{FileLoader, SourceDatabase};
    // hack until we get hygiene working (same character amount to preserve formatting as much as possible)
    const DOLLAR_CRATE_REPLACE: &str = &"__r_a_";
    let expansion = expansion.replace("$crate", DOLLAR_CRATE_REPLACE);
    let (prefix, suffix) = match kind {
        SyntaxKind::MACRO_PAT => ("fn __(", ": u32);"),
        SyntaxKind::MACRO_EXPR | SyntaxKind::MACRO_STMTS => ("fn __() {", "}"),
        SyntaxKind::MACRO_TYPE => ("type __ =", ";"),
        _ => ("", ""),
    };
    let expansion = format!("{prefix}{expansion}{suffix}");

    let &crate_id = db.relevant_crates(file_id).iter().next()?;
    let edition = db.crate_graph()[crate_id].edition;

    let mut cmd = std::process::Command::new(toolchain::rustfmt());
    cmd.arg("--edition");
    cmd.arg(edition.to_string());

    let mut rustfmt = cmd
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .spawn()
        .ok()?;

    std::io::Write::write_all(&mut rustfmt.stdin.as_mut()?, expansion.as_bytes()).ok()?;

    let output = rustfmt.wait_with_output().ok()?;
    let captured_stdout = String::from_utf8(output.stdout).ok()?;

    if output.status.success() && !captured_stdout.trim().is_empty() {
        let output = captured_stdout.replace(DOLLAR_CRATE_REPLACE, "$crate");
        let output = output.trim().strip_prefix(prefix)?;
        let output = match kind {
            SyntaxKind::MACRO_PAT => {
                output.strip_suffix(suffix).or_else(|| output.strip_suffix(": u32,\n);"))?
            }
            _ => output.strip_suffix(suffix)?,
        };
        let trim_indent = stdx::trim_indent(output);
        tracing::debug!("expand_macro: formatting succeeded");
        Some(trim_indent)
    } else {
        None
    }
}

#[cfg(test)]
mod tests {
    use expect_test::{expect, Expect};

    use crate::fixture;

    #[track_caller]
    fn check(ra_fixture: &str, expect: Expect) {
        let (analysis, pos) = fixture::position(ra_fixture);
        let expansion = analysis.expand_macro(pos).unwrap().unwrap();
        let actual = format!("{}\n{}", expansion.name, expansion.expansion);
        expect.assert_eq(&actual);
    }

    #[test]
    fn macro_expand_as_keyword() {
        check(
            r#"
macro_rules! bar {
    ($i:tt) => { $i as _ }
}
fn main() {
    let x: u64 = ba$0r!(5i64);
}
"#,
            expect![[r#"
                bar
                5i64 as _"#]],
        );
    }

    #[test]
    fn macro_expand_underscore() {
        check(
            r#"
macro_rules! bar {
    ($i:tt) => { for _ in 0..$i {} }
}
fn main() {
    ba$0r!(42);
}
"#,
            expect![[r#"
                bar
                for _ in 0..42{}
                "#]],
        );
    }

    #[test]
    fn macro_expand_recursive_expansion() {
        check(
            r#"
macro_rules! bar {
    () => { fn  b() {} }
}
macro_rules! foo {
    () => { bar!(); }
}
macro_rules! baz {
    () => { foo!(); }
}
f$0oo!();
"#,
            expect![[r#"
                foo
                fn b(){}

            "#]],
        );
    }

    #[test]
    fn macro_expand_multiple_lines() {
        check(
            r#"
macro_rules! foo {
    () => {
        fn some_thing() -> u32 {
            let a = 0;
            a + 10
        }
    }
}
f$0oo!();
        "#,
            expect![[r#"
                foo
                fn some_thing() -> u32 {
                  let a = 0;
                  a+10
                }
            "#]],
        );
    }

    #[test]
    fn macro_expand_match_ast() {
        check(
            r#"
macro_rules! match_ast {
    (match $node:ident { $($tt:tt)* }) => { match_ast!(match ($node) { $($tt)* }) };
    (match ($node:expr) {
        $( ast::$ast:ident($it:ident) => $res:block, )*
        _ => $catch_all:expr $(,)?
    }) => {{
        $( if let Some($it) = ast::$ast::cast($node.clone()) $res else )*
        { $catch_all }
    }};
}

fn main() {
    mat$0ch_ast! {
        match container {
            ast::TraitDef(it) => {},
            ast::ImplDef(it) => {},
            _ => { continue },
        }
    }
}
"#,
            expect![[r#"
       match_ast
       {
         if let Some(it) = ast::TraitDef::cast(container.clone()){}
         else if let Some(it) = ast::ImplDef::cast(container.clone()){}
         else {
           {
             continue
           }
         }
       }"#]],
        );
    }

    #[test]
    fn macro_expand_match_ast_inside_let_statement() {
        check(
            r#"
macro_rules! match_ast {
    (match $node:ident { $($tt:tt)* }) => { match_ast!(match ($node) { $($tt)* }) };
    (match ($node:expr) {}) => {{}};
}

fn main() {
    let p = f(|it| {
        let res = mat$0ch_ast! { match c {}};
        Some(res)
    })?;
}
"#,
            expect![[r#"
                match_ast
                {}
            "#]],
        );
    }

    #[test]
    fn macro_expand_inner_macro_fail_to_expand() {
        check(
            r#"
macro_rules! bar {
    (BAD) => {};
}
macro_rules! foo {
    () => {bar!()};
}

fn main() {
    let res = fo$0o!();
}
"#,
            expect![[r#"
                foo
            "#]],
        );
    }

    #[test]
    fn macro_expand_with_dollar_crate() {
        check(
            r#"
#[macro_export]
macro_rules! bar {
    () => {0};
}
macro_rules! foo {
    () => {$crate::bar!()};
}

fn main() {
    let res = fo$0o!();
}
"#,
            expect![[r#"
                foo
                0"#]],
        );
    }

    #[test]
    fn macro_expand_with_dyn_absolute_path() {
        check(
            r#"
macro_rules! foo {
    () => {fn f<T>(_: &dyn ::std::marker::Copy) {}};
}

fn main() {
    let res = fo$0o!();
}
"#,
            expect![[r#"
                foo
                fn f<T>(_: &dyn ::std::marker::Copy){}
            "#]],
        );
    }

    #[test]
    fn macro_expand_derive() {
        check(
            r#"
//- proc_macros: identity
//- minicore: clone, derive

#[proc_macros::identity]
#[derive(C$0lone)]
struct Foo {}
"#,
            expect![[r#"
                Clone
                impl < >core::clone::Clone for Foo< >{}

            "#]],
        );
    }

    #[test]
    fn macro_expand_derive2() {
        check(
            r#"
//- minicore: copy, clone, derive

#[derive(Cop$0y)]
#[derive(Clone)]
struct Foo {}
"#,
            expect![[r#"
                Copy
                impl < >core::marker::Copy for Foo< >{}

            "#]],
        );
    }

    #[test]
    fn macro_expand_derive_multi() {
        check(
            r#"
//- minicore: copy, clone, derive

#[derive(Cop$0y, Clone)]
struct Foo {}
"#,
            expect![[r#"
                Copy
                impl < >core::marker::Copy for Foo< >{}

            "#]],
        );
        check(
            r#"
//- minicore: copy, clone, derive

#[derive(Copy, Cl$0one)]
struct Foo {}
"#,
            expect![[r#"
                Clone
                impl < >core::clone::Clone for Foo< >{}

            "#]],
        );
    }
}
<|endoftext|>


<commit_before>//! Lexing `&str` into a sequence of Rust tokens.
//!
//! Note that strictly speaking the parser in this crate is not required to work
//! on tokens which originated from text. Macros, eg, can synthesize tokes out
//! of thin air. So, ideally, lexer should be an orthogonal crate. It is however
//! convenient to include a text-based lexer here!
//!
//! Note that these tokens, unlike the tokens we feed into the parser, do
//! include info about comments and whitespace.

use crate::{
    SyntaxKind::{self, *},
    T,
};

pub struct LexedStr<'a> {
    text: &'a str,
    kind: Vec<SyntaxKind>,
    start: Vec<u32>,
    error: Vec<LexError>,
}

struct LexError {
    msg: String,
    token: u32,
}

impl<'a> LexedStr<'a> {
    pub fn new(text: &'a str) -> LexedStr<'a> {
        let mut res = LexedStr { text, kind: Vec::new(), start: Vec::new(), error: Vec::new() };

        let mut offset = 0;
        if let Some(shebang_len) = rustc_lexer::strip_shebang(text) {
            res.push(SHEBANG, offset);
            offset = shebang_len
        };
        for token in rustc_lexer::tokenize(&text[offset..]) {
            let token_text = &text[offset..][..token.len];

            let (kind, err) = from_rustc(&token.kind, token_text);
            res.push(kind, offset);
            offset += token.len;

            if let Some(err) = err {
                let token = res.len() as u32;
                let msg = err.to_string();
                res.error.push(LexError { msg, token });
            }
        }
        res.push(EOF, offset);

        res
    }

    pub fn single_token(text: &'a str) -> Option<SyntaxKind> {
        if text.is_empty() {
            return None;
        }

        let token = rustc_lexer::first_token(text);
        if token.len != text.len() {
            return None;
        }

        let (kind, err) = from_rustc(&token.kind, text);
        if err.is_some() {
            return None;
        }

        Some(kind)
    }

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

    pub fn len(&self) -> usize {
        self.kind.len() - 1
    }

    pub fn kind(&self, i: usize) -> SyntaxKind {
        assert!(i < self.len());
        self.kind[i]
    }
    pub fn text(&self, i: usize) -> &str {
        assert!(i < self.len());
        let lo = self.start[i] as usize;
        let hi = self.start[i + 1] as usize;
        &self.text[lo..hi]
    }
    pub fn error(&self, i: usize) -> Option<&str> {
        assert!(i < self.len());
        let err = self.error.binary_search_by_key(&(i as u32), |i| i.token).ok()?;
        Some(self.error[err].msg.as_str())
    }

    fn push(&mut self, kind: SyntaxKind, offset: usize) {
        self.kind.push(kind);
        self.start.push(offset as u32);
    }
}

/// Returns `SyntaxKind` and an optional tokenize error message.
fn from_rustc(
    kind: &rustc_lexer::TokenKind,
    token_text: &str,
) -> (SyntaxKind, Option<&'static str>) {
    // A note on an intended tradeoff:
    // We drop some useful information here (see patterns with double dots `..`)
    // Storing that info in `SyntaxKind` is not possible due to its layout requirements of
    // being `u16` that come from `rowan::SyntaxKind`.
    let mut err = "";

    let syntax_kind = {
        match kind {
            rustc_lexer::TokenKind::LineComment { doc_style: _ } => COMMENT,
            rustc_lexer::TokenKind::BlockComment { doc_style: _, terminated } => {
                if !terminated {
                    err = "Missing trailing `*/` symbols to terminate the block comment";
                }
                COMMENT
            }

            rustc_lexer::TokenKind::Whitespace => WHITESPACE,

            rustc_lexer::TokenKind::Ident if token_text == "_" => UNDERSCORE,
            rustc_lexer::TokenKind::Ident => SyntaxKind::from_keyword(token_text).unwrap_or(IDENT),

            rustc_lexer::TokenKind::RawIdent => IDENT,
            rustc_lexer::TokenKind::Literal { kind, .. } => return from_rustc_literal(kind),

            rustc_lexer::TokenKind::Lifetime { starts_with_number } => {
                if *starts_with_number {
                    err = "Lifetime name cannot start with a number";
                }
                LIFETIME_IDENT
            }

            rustc_lexer::TokenKind::Semi => T![;],
            rustc_lexer::TokenKind::Comma => T![,],
            rustc_lexer::TokenKind::Dot => T![.],
            rustc_lexer::TokenKind::OpenParen => T!['('],
            rustc_lexer::TokenKind::CloseParen => T![')'],
            rustc_lexer::TokenKind::OpenBrace => T!['{'],
            rustc_lexer::TokenKind::CloseBrace => T!['}'],
            rustc_lexer::TokenKind::OpenBracket => T!['['],
            rustc_lexer::TokenKind::CloseBracket => T![']'],
            rustc_lexer::TokenKind::At => T![@],
            rustc_lexer::TokenKind::Pound => T![#],
            rustc_lexer::TokenKind::Tilde => T![~],
            rustc_lexer::TokenKind::Question => T![?],
            rustc_lexer::TokenKind::Colon => T![:],
            rustc_lexer::TokenKind::Dollar => T![$],
            rustc_lexer::TokenKind::Eq => T![=],
            rustc_lexer::TokenKind::Bang => T![!],
            rustc_lexer::TokenKind::Lt => T![<],
            rustc_lexer::TokenKind::Gt => T![>],
            rustc_lexer::TokenKind::Minus => T![-],
            rustc_lexer::TokenKind::And => T![&],
            rustc_lexer::TokenKind::Or => T![|],
            rustc_lexer::TokenKind::Plus => T![+],
            rustc_lexer::TokenKind::Star => T![*],
            rustc_lexer::TokenKind::Slash => T![/],
            rustc_lexer::TokenKind::Caret => T![^],
            rustc_lexer::TokenKind::Percent => T![%],
            rustc_lexer::TokenKind::Unknown => ERROR,
        }
    };

    let err = if err.is_empty() { None } else { Some(err) };
    (syntax_kind, err)
}

fn from_rustc_literal(kind: &rustc_lexer::LiteralKind) -> (SyntaxKind, Option<&'static str>) {
    let mut err = "";

    let syntax_kind = match *kind {
        rustc_lexer::LiteralKind::Int { empty_int, base: _ } => {
            if empty_int {
                err = "Missing digits after the integer base prefix";
            }
            INT_NUMBER
        }
        rustc_lexer::LiteralKind::Float { empty_exponent, base: _ } => {
            if empty_exponent {
                err = "Missing digits after the exponent symbol";
            }
            FLOAT_NUMBER
        }
        rustc_lexer::LiteralKind::Char { terminated } => {
            if !terminated {
                err = "Missing trailing `'` symbol to terminate the character literal";
            }
            CHAR
        }
        rustc_lexer::LiteralKind::Byte { terminated } => {
            if !terminated {
                err = "Missing trailing `'` symbol to terminate the byte literal";
            }
            BYTE
        }
        rustc_lexer::LiteralKind::Str { terminated } => {
            if !terminated {
                err = "Missing trailing `\"` symbol to terminate the string literal";
            }
            STRING
        }
        rustc_lexer::LiteralKind::ByteStr { terminated } => {
            if !terminated {
                err = "Missing trailing `\"` symbol to terminate the byte string literal";
            }
            BYTE_STRING
        }
        rustc_lexer::LiteralKind::RawStr { err: raw_str_err, .. } => {
            if let Some(raw_str_err) = raw_str_err {
                err = match raw_str_err {
                    rustc_lexer::RawStrError::InvalidStarter { .. } => "Missing `\"` symbol after `#` symbols to begin the raw string literal",
                    rustc_lexer::RawStrError::NoTerminator { expected, found, .. } => if expected == found {
                        "Missing trailing `\"` to terminate the raw string literal"
                    } else {
                        "Missing trailing `\"` with `#` symbols to terminate the raw string literal"
                    },
                    rustc_lexer::RawStrError::TooManyDelimiters { .. } => "Too many `#` symbols: raw strings may be delimited by up to 65535 `#` symbols",
                };
            };
            STRING
        }
        rustc_lexer::LiteralKind::RawByteStr { err: raw_str_err, .. } => {
            if let Some(raw_str_err) = raw_str_err {
                err = match raw_str_err {
                    rustc_lexer::RawStrError::InvalidStarter { .. } => "Missing `\"` symbol after `#` symbols to begin the raw byte string literal",
                    rustc_lexer::RawStrError::NoTerminator { expected, found, .. } => if expected == found {
                        "Missing trailing `\"` to terminate the raw byte string literal"
                    } else {
                        "Missing trailing `\"` with `#` symbols to terminate the raw byte string literal"
                    },
                    rustc_lexer::RawStrError::TooManyDelimiters { .. } => "Too many `#` symbols: raw byte strings may be delimited by up to 65535 `#` symbols",
                };
            };

            BYTE_STRING
        }
    };

    let err = if err.is_empty() { None } else { Some(err) };
    (syntax_kind, err)
}
<commit_msg>converting lexed str to tokens<commit_after>//! Lexing `&str` into a sequence of Rust tokens.
//!
//! Note that strictly speaking the parser in this crate is not required to work
//! on tokens which originated from text. Macros, eg, can synthesize tokes out
//! of thin air. So, ideally, lexer should be an orthogonal crate. It is however
//! convenient to include a text-based lexer here!
//!
//! Note that these tokens, unlike the tokens we feed into the parser, do
//! include info about comments and whitespace.

use crate::{
    SyntaxKind::{self, *},
    T,
};

pub struct LexedStr<'a> {
    text: &'a str,
    kind: Vec<SyntaxKind>,
    start: Vec<u32>,
    error: Vec<LexError>,
}

struct LexError {
    msg: String,
    token: u32,
}

impl<'a> LexedStr<'a> {
    pub fn new(text: &'a str) -> LexedStr<'a> {
        let mut res = LexedStr { text, kind: Vec::new(), start: Vec::new(), error: Vec::new() };

        let mut offset = 0;
        if let Some(shebang_len) = rustc_lexer::strip_shebang(text) {
            res.push(SHEBANG, offset);
            offset = shebang_len
        };
        for token in rustc_lexer::tokenize(&text[offset..]) {
            let token_text = &text[offset..][..token.len];

            let (kind, err) = from_rustc(&token.kind, token_text);
            res.push(kind, offset);
            offset += token.len;

            if let Some(err) = err {
                let token = res.len() as u32;
                let msg = err.to_string();
                res.error.push(LexError { msg, token });
            }
        }
        res.push(EOF, offset);

        res
    }

    pub fn single_token(text: &'a str) -> Option<SyntaxKind> {
        if text.is_empty() {
            return None;
        }

        let token = rustc_lexer::first_token(text);
        if token.len != text.len() {
            return None;
        }

        let (kind, err) = from_rustc(&token.kind, text);
        if err.is_some() {
            return None;
        }

        Some(kind)
    }

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

    pub fn len(&self) -> usize {
        self.kind.len() - 1
    }

    pub fn kind(&self, i: usize) -> SyntaxKind {
        assert!(i < self.len());
        self.kind[i]
    }

    pub fn text(&self, i: usize) -> &str {
        assert!(i < self.len());
        let lo = self.start[i] as usize;
        let hi = self.start[i + 1] as usize;
        &self.text[lo..hi]
    }

    pub fn error(&self, i: usize) -> Option<&str> {
        assert!(i < self.len());
        let err = self.error.binary_search_by_key(&(i as u32), |i| i.token).ok()?;
        Some(self.error[err].msg.as_str())
    }

    pub fn to_tokens(&self) -> crate::Tokens {
        let mut res = crate::Tokens::default();
        let mut was_joint = false;
        for i in 0..self.len() {
            let kind = self.kind(i);
            if kind.is_trivia() {
                was_joint = false
            } else {
                if kind == SyntaxKind::IDENT {
                    let token_text = self.text(i);
                    let contextual_kw = SyntaxKind::from_contextual_keyword(token_text)
                        .unwrap_or(SyntaxKind::IDENT);
                    res.push_ident(contextual_kw);
                } else {
                    if was_joint {
                        res.was_joint();
                    }
                    res.push(kind);
                }
                was_joint = true;
            }
        }
        res
    }

    fn push(&mut self, kind: SyntaxKind, offset: usize) {
        self.kind.push(kind);
        self.start.push(offset as u32);
    }
}

/// Returns `SyntaxKind` and an optional tokenize error message.
fn from_rustc(
    kind: &rustc_lexer::TokenKind,
    token_text: &str,
) -> (SyntaxKind, Option<&'static str>) {
    // A note on an intended tradeoff:
    // We drop some useful information here (see patterns with double dots `..`)
    // Storing that info in `SyntaxKind` is not possible due to its layout requirements of
    // being `u16` that come from `rowan::SyntaxKind`.
    let mut err = "";

    let syntax_kind = {
        match kind {
            rustc_lexer::TokenKind::LineComment { doc_style: _ } => COMMENT,
            rustc_lexer::TokenKind::BlockComment { doc_style: _, terminated } => {
                if !terminated {
                    err = "Missing trailing `*/` symbols to terminate the block comment";
                }
                COMMENT
            }

            rustc_lexer::TokenKind::Whitespace => WHITESPACE,

            rustc_lexer::TokenKind::Ident if token_text == "_" => UNDERSCORE,
            rustc_lexer::TokenKind::Ident => SyntaxKind::from_keyword(token_text).unwrap_or(IDENT),

            rustc_lexer::TokenKind::RawIdent => IDENT,
            rustc_lexer::TokenKind::Literal { kind, .. } => return from_rustc_literal(kind),

            rustc_lexer::TokenKind::Lifetime { starts_with_number } => {
                if *starts_with_number {
                    err = "Lifetime name cannot start with a number";
                }
                LIFETIME_IDENT
            }

            rustc_lexer::TokenKind::Semi => T![;],
            rustc_lexer::TokenKind::Comma => T![,],
            rustc_lexer::TokenKind::Dot => T![.],
            rustc_lexer::TokenKind::OpenParen => T!['('],
            rustc_lexer::TokenKind::CloseParen => T![')'],
            rustc_lexer::TokenKind::OpenBrace => T!['{'],
            rustc_lexer::TokenKind::CloseBrace => T!['}'],
            rustc_lexer::TokenKind::OpenBracket => T!['['],
            rustc_lexer::TokenKind::CloseBracket => T![']'],
            rustc_lexer::TokenKind::At => T![@],
            rustc_lexer::TokenKind::Pound => T![#],
            rustc_lexer::TokenKind::Tilde => T![~],
            rustc_lexer::TokenKind::Question => T![?],
            rustc_lexer::TokenKind::Colon => T![:],
            rustc_lexer::TokenKind::Dollar => T![$],
            rustc_lexer::TokenKind::Eq => T![=],
            rustc_lexer::TokenKind::Bang => T![!],
            rustc_lexer::TokenKind::Lt => T![<],
            rustc_lexer::TokenKind::Gt => T![>],
            rustc_lexer::TokenKind::Minus => T![-],
            rustc_lexer::TokenKind::And => T![&],
            rustc_lexer::TokenKind::Or => T![|],
            rustc_lexer::TokenKind::Plus => T![+],
            rustc_lexer::TokenKind::Star => T![*],
            rustc_lexer::TokenKind::Slash => T![/],
            rustc_lexer::TokenKind::Caret => T![^],
            rustc_lexer::TokenKind::Percent => T![%],
            rustc_lexer::TokenKind::Unknown => ERROR,
        }
    };

    let err = if err.is_empty() { None } else { Some(err) };
    (syntax_kind, err)
}

fn from_rustc_literal(kind: &rustc_lexer::LiteralKind) -> (SyntaxKind, Option<&'static str>) {
    let mut err = "";

    let syntax_kind = match *kind {
        rustc_lexer::LiteralKind::Int { empty_int, base: _ } => {
            if empty_int {
                err = "Missing digits after the integer base prefix";
            }
            INT_NUMBER
        }
        rustc_lexer::LiteralKind::Float { empty_exponent, base: _ } => {
            if empty_exponent {
                err = "Missing digits after the exponent symbol";
            }
            FLOAT_NUMBER
        }
        rustc_lexer::LiteralKind::Char { terminated } => {
            if !terminated {
                err = "Missing trailing `'` symbol to terminate the character literal";
            }
            CHAR
        }
        rustc_lexer::LiteralKind::Byte { terminated } => {
            if !terminated {
                err = "Missing trailing `'` symbol to terminate the byte literal";
            }
            BYTE
        }
        rustc_lexer::LiteralKind::Str { terminated } => {
            if !terminated {
                err = "Missing trailing `\"` symbol to terminate the string literal";
            }
            STRING
        }
        rustc_lexer::LiteralKind::ByteStr { terminated } => {
            if !terminated {
                err = "Missing trailing `\"` symbol to terminate the byte string literal";
            }
            BYTE_STRING
        }
        rustc_lexer::LiteralKind::RawStr { err: raw_str_err, .. } => {
            if let Some(raw_str_err) = raw_str_err {
                err = match raw_str_err {
                    rustc_lexer::RawStrError::InvalidStarter { .. } => "Missing `\"` symbol after `#` symbols to begin the raw string literal",
                    rustc_lexer::RawStrError::NoTerminator { expected, found, .. } => if expected == found {
                        "Missing trailing `\"` to terminate the raw string literal"
                    } else {
                        "Missing trailing `\"` with `#` symbols to terminate the raw string literal"
                    },
                    rustc_lexer::RawStrError::TooManyDelimiters { .. } => "Too many `#` symbols: raw strings may be delimited by up to 65535 `#` symbols",
                };
            };
            STRING
        }
        rustc_lexer::LiteralKind::RawByteStr { err: raw_str_err, .. } => {
            if let Some(raw_str_err) = raw_str_err {
                err = match raw_str_err {
                    rustc_lexer::RawStrError::InvalidStarter { .. } => "Missing `\"` symbol after `#` symbols to begin the raw byte string literal",
                    rustc_lexer::RawStrError::NoTerminator { expected, found, .. } => if expected == found {
                        "Missing trailing `\"` to terminate the raw byte string literal"
                    } else {
                        "Missing trailing `\"` with `#` symbols to terminate the raw byte string literal"
                    },
                    rustc_lexer::RawStrError::TooManyDelimiters { .. } => "Too many `#` symbols: raw byte strings may be delimited by up to 65535 `#` symbols",
                };
            };

            BYTE_STRING
        }
    };

    let err = if err.is_empty() { None } else { Some(err) };
    (syntax_kind, err)
}
<|endoftext|>


<commit_before>use rustup_utils::{utils};
use errors::*;
use temp;
use notifications::*;
use sha2::{Sha256, Digest};
use url::Url;

use std::path::{Path, PathBuf};
use std::fs;
use std::ops;

const UPDATE_HASH_LEN: usize = 20;

#[derive(Copy, Clone)]
pub struct DownloadCfg<'a> {
    pub dist_root: &'a str,
    pub temp_cfg: &'a temp::Cfg,
    pub download_dir: &'a PathBuf,
    pub notify_handler: &'a Fn(Notification),
}

pub struct File {
    path: PathBuf,
}

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

    fn deref(&self) -> &Path {
        ops::Deref::deref(&self.path)
    }
}

impl<'a> DownloadCfg<'a> {

    /// Downloads a file, validating its hash, and resuming interrupted downloads
    /// Partial downloads are stored in `self.download_dir`, keyed by hash. If the
    /// target file already exists, then the hash is checked and it is returned
    /// immediately without re-downloading.
    pub fn download(&self, url: &Url, hash: &str) -> Result<File> {

        try!(utils::ensure_dir_exists("Download Directory", &self.download_dir, &|n| (self.notify_handler)(n.into())));
        let target_file = self.download_dir.join(Path::new(hash));

        if target_file.exists() {
            let cached_result = try!(file_hash(&target_file));
            if hash == cached_result {
                (self.notify_handler)(Notification::FileAlreadyDownloaded);
                (self.notify_handler)(Notification::ChecksumValid(&url.to_string()));
                return Ok(File { path: target_file, });
            } else {
                (self.notify_handler)(Notification::CachedFileChecksumFailed);
                try!(fs::remove_file(&target_file).chain_err(|| "cleaning up previous download"));
            }
        }

        let partial_file_path =
            target_file.with_file_name(
                target_file.file_name().map(|s| {
                    s.to_str().unwrap_or("_")})
                    .unwrap_or("_")
                    .to_owned()
                + ".partial");

        let mut hasher = Sha256::new();

        try!(utils::download_file_with_resume(&url,
                                  &partial_file_path,
                                  Some(&mut hasher),
                                  true,
                                  &|n| (self.notify_handler)(n.into())));

        let actual_hash = hasher.result_str();

        if hash != actual_hash {
            // Incorrect hash
            return Err(ErrorKind::ChecksumFailed {
                url: url.to_string(),
                expected: hash.to_string(),
                calculated: actual_hash,
            }.into());
        } else {
            (self.notify_handler)(Notification::ChecksumValid(&url.to_string()));
            try!(fs::rename(&partial_file_path, &target_file));
            return Ok(File { path: target_file });
        }
    }

    pub fn clean(&self, hashes: &Vec<String>) -> Result<()> {
        for hash in hashes.iter() {
            let used_file = self.download_dir.join(hash);
            if self.download_dir.join(&used_file).exists() {
                try!(fs::remove_file(used_file).chain_err(|| "cleaning up cached downloads"));
            }
        }
        Ok(())
    }

    fn download_hash(&self, url: &str) -> Result<String> {
        let hash_url = try!(utils::parse_url(&(url.to_owned() + ".sha256")));
        let hash_file = try!(self.temp_cfg.new_file());

        try!(utils::download_file(&hash_url,
                                &hash_file,
                                None,
                                &|n| (self.notify_handler)(n.into())));

        Ok(try!(utils::read_file("hash", &hash_file).map(|s| s[0..64].to_owned())))
    }

    /// Downloads a file, sourcing its hash from the same url with a `.sha256` suffix.
    /// If `update_hash` is present, then that will be compared to the downloaded hash,
    /// and if they match, the download is skipped.
    pub fn download_and_check(&self,
                                url_str: &str,
                                update_hash: Option<&Path>,
                                ext: &str)
                                -> Result<Option<(temp::File<'a>, String)>> {
        let hash = try!(self.download_hash(url_str));
        let partial_hash: String = hash.chars().take(UPDATE_HASH_LEN).collect();

        if let Some(hash_file) = update_hash {

            if utils::is_file(hash_file) {
                if let Ok(contents) = utils::read_file("update hash", hash_file) {
                    if contents == partial_hash {
                        // Skip download, update hash matches
                        return Ok(None);
                    }
                } else {
                    (self.notify_handler)(Notification::CantReadUpdateHash(hash_file));
                }
            } else {
                (self.notify_handler)(Notification::NoUpdateHash(hash_file));
            }
        }

        let url = try!(utils::parse_url(url_str));
        let file = try!(self.temp_cfg.new_file_with_ext("", ext));

        let mut hasher = Sha256::new();
        try!(utils::download_file(&url,
                                &file,
                                Some(&mut hasher),
                                &|n| (self.notify_handler)(n.into())));
        let actual_hash = hasher.result_str();

        if hash != actual_hash {
            // Incorrect hash
            return Err(ErrorKind::ChecksumFailed {
                    url: url_str.to_owned(),
                    expected: hash,
                    calculated: actual_hash,
                }
                .into());
        } else {
            (self.notify_handler)(Notification::ChecksumValid(url_str));
        }

        // TODO: Check the signature of the file

        Ok(Some((file, partial_hash)))
    }
}

fn file_hash(path: &Path) -> Result<String> {
    let mut hasher = Sha256::new();
    use std::io::Read;
    let mut downloaded = try!(fs::File::open(&path).chain_err(|| "opening already downloaded file"));
    let mut buf = vec![0; 32768];
    loop {
        if let Ok(n) = downloaded.read(&mut buf) {
            if n == 0 { break; }
            hasher.input(&buf[..n]);
        } else {
            break;
        }
    }

    Ok(hasher.result_str())
}<commit_msg>Add missing newline<commit_after>use rustup_utils::{utils};
use errors::*;
use temp;
use notifications::*;
use sha2::{Sha256, Digest};
use url::Url;

use std::path::{Path, PathBuf};
use std::fs;
use std::ops;

const UPDATE_HASH_LEN: usize = 20;

#[derive(Copy, Clone)]
pub struct DownloadCfg<'a> {
    pub dist_root: &'a str,
    pub temp_cfg: &'a temp::Cfg,
    pub download_dir: &'a PathBuf,
    pub notify_handler: &'a Fn(Notification),
}

pub struct File {
    path: PathBuf,
}

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

    fn deref(&self) -> &Path {
        ops::Deref::deref(&self.path)
    }
}

impl<'a> DownloadCfg<'a> {

    /// Downloads a file, validating its hash, and resuming interrupted downloads
    /// Partial downloads are stored in `self.download_dir`, keyed by hash. If the
    /// target file already exists, then the hash is checked and it is returned
    /// immediately without re-downloading.
    pub fn download(&self, url: &Url, hash: &str) -> Result<File> {

        try!(utils::ensure_dir_exists("Download Directory", &self.download_dir, &|n| (self.notify_handler)(n.into())));
        let target_file = self.download_dir.join(Path::new(hash));

        if target_file.exists() {
            let cached_result = try!(file_hash(&target_file));
            if hash == cached_result {
                (self.notify_handler)(Notification::FileAlreadyDownloaded);
                (self.notify_handler)(Notification::ChecksumValid(&url.to_string()));
                return Ok(File { path: target_file, });
            } else {
                (self.notify_handler)(Notification::CachedFileChecksumFailed);
                try!(fs::remove_file(&target_file).chain_err(|| "cleaning up previous download"));
            }
        }

        let partial_file_path =
            target_file.with_file_name(
                target_file.file_name().map(|s| {
                    s.to_str().unwrap_or("_")})
                    .unwrap_or("_")
                    .to_owned()
                + ".partial");

        let mut hasher = Sha256::new();

        try!(utils::download_file_with_resume(&url,
                                  &partial_file_path,
                                  Some(&mut hasher),
                                  true,
                                  &|n| (self.notify_handler)(n.into())));

        let actual_hash = hasher.result_str();

        if hash != actual_hash {
            // Incorrect hash
            return Err(ErrorKind::ChecksumFailed {
                url: url.to_string(),
                expected: hash.to_string(),
                calculated: actual_hash,
            }.into());
        } else {
            (self.notify_handler)(Notification::ChecksumValid(&url.to_string()));
            try!(fs::rename(&partial_file_path, &target_file));
            return Ok(File { path: target_file });
        }
    }

    pub fn clean(&self, hashes: &Vec<String>) -> Result<()> {
        for hash in hashes.iter() {
            let used_file = self.download_dir.join(hash);
            if self.download_dir.join(&used_file).exists() {
                try!(fs::remove_file(used_file).chain_err(|| "cleaning up cached downloads"));
            }
        }
        Ok(())
    }

    fn download_hash(&self, url: &str) -> Result<String> {
        let hash_url = try!(utils::parse_url(&(url.to_owned() + ".sha256")));
        let hash_file = try!(self.temp_cfg.new_file());

        try!(utils::download_file(&hash_url,
                                &hash_file,
                                None,
                                &|n| (self.notify_handler)(n.into())));

        Ok(try!(utils::read_file("hash", &hash_file).map(|s| s[0..64].to_owned())))
    }

    /// Downloads a file, sourcing its hash from the same url with a `.sha256` suffix.
    /// If `update_hash` is present, then that will be compared to the downloaded hash,
    /// and if they match, the download is skipped.
    pub fn download_and_check(&self,
                                url_str: &str,
                                update_hash: Option<&Path>,
                                ext: &str)
                                -> Result<Option<(temp::File<'a>, String)>> {
        let hash = try!(self.download_hash(url_str));
        let partial_hash: String = hash.chars().take(UPDATE_HASH_LEN).collect();

        if let Some(hash_file) = update_hash {

            if utils::is_file(hash_file) {
                if let Ok(contents) = utils::read_file("update hash", hash_file) {
                    if contents == partial_hash {
                        // Skip download, update hash matches
                        return Ok(None);
                    }
                } else {
                    (self.notify_handler)(Notification::CantReadUpdateHash(hash_file));
                }
            } else {
                (self.notify_handler)(Notification::NoUpdateHash(hash_file));
            }
        }

        let url = try!(utils::parse_url(url_str));
        let file = try!(self.temp_cfg.new_file_with_ext("", ext));

        let mut hasher = Sha256::new();
        try!(utils::download_file(&url,
                                &file,
                                Some(&mut hasher),
                                &|n| (self.notify_handler)(n.into())));
        let actual_hash = hasher.result_str();

        if hash != actual_hash {
            // Incorrect hash
            return Err(ErrorKind::ChecksumFailed {
                    url: url_str.to_owned(),
                    expected: hash,
                    calculated: actual_hash,
                }
                .into());
        } else {
            (self.notify_handler)(Notification::ChecksumValid(url_str));
        }

        // TODO: Check the signature of the file

        Ok(Some((file, partial_hash)))
    }
}

fn file_hash(path: &Path) -> Result<String> {
    let mut hasher = Sha256::new();
    use std::io::Read;
    let mut downloaded = try!(fs::File::open(&path).chain_err(|| "opening already downloaded file"));
    let mut buf = vec![0; 32768];
    loop {
        if let Ok(n) = downloaded.read(&mut buf) {
            if n == 0 { break; }
            hasher.input(&buf[..n]);
        } else {
            break;
        }
    }

    Ok(hasher.result_str())
}
<|endoftext|>


<commit_before>extern crate clap;

use clap::Parser;

use std::collections::{HashMap, HashSet};
use std::fmt;
use std::time::Instant;

use cubes_rs::{Piece, Puzzle};

const SIN: [i32; 4] = [0, 1, 0, -1];
const COS: [i32; 4] = [1, 0, -1, 0];

/// Program to display the 3x3 cube solution(s).
/// Solves mintaur by default.
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// Solve the blue problem.
    #[arg(short, long)]
    blue: bool,

    /// Solve the green problem.
    #[arg(short, long)]
    green: bool,
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
struct PuzzleDense {
    data: [[[i32; 3]; 3]; 3],
}

impl fmt::Display for PuzzleDense {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "{:?}", self.data[0])?;
        writeln!(f, "{:?}", self.data[1])?;
        writeln!(f, "{:?}", self.data[2])
    }
}

fn main() {
    let now = Instant::now();

    let puzzle;
    let args = Args::parse();
    if args.blue {
        puzzle = cubes_rs::blue();
    } else if args.green {
        puzzle = cubes_rs::green();
    } else {
        puzzle = cubes_rs::minotaur();
    }

    let pieces = push_to_zero(puzzle);
    for piece in &pieces {
        for part in piece {
            println!("{:?}", part);
        }
        println!();
    }

    let zero = vec![PuzzleDense { data: zeros() }];

    let mut solutions = Vec::new();
    for (i, piece) in pieces.iter().enumerate() {
        if i == 0 {
            solutions = all_puts(zero.clone(), (i + 1) as i32, piece);
        } else {
            solutions = all_rotations_and_puts(solutions.clone(), (i + 1) as i32, piece);
        }
    }

    let solutions = unique_pieces(solutions);
    for solution in solutions {
        println!("{:}", solution);
    }

    let elapsed_time = now.elapsed();
    println!("Running the program took {} seconds.", elapsed_time.as_millis() as f64 / 1000.0);
}

fn unique_pieces(puzzles: Vec<PuzzleDense>) -> Vec<PuzzleDense> {
    let mut puzzles_unique = HashSet::new();
    for mut puzzle in puzzles {
        let mut puzzle_unique: HashMap<i32, i32> = HashMap::new();
        let mut piece_count = 0;
        for x in 0..3 {
            for y in 0..3 {
                for z in 0..3 {
                    let value = puzzle.data[x][y][z];
                    if puzzle_unique.contains_key(&value) {
                        puzzle.data[x][y][z] = puzzle_unique[&value]
                    } else {
                        piece_count += 1;
                        puzzle_unique.insert(value, piece_count);
                        puzzle.data[x][y][z] = puzzle_unique[&value]
                    }
                }
            }
        }
        puzzles_unique.insert(puzzle);
    }

    let mut all_puzzles = Vec::new();
    for puzzle in puzzles_unique {
        all_puzzles.push(puzzle);
    }
    all_puzzles
}

fn all_rotations_and_puts(
    already_placed: Vec<PuzzleDense>,
    piece_count: i32,
    piece: &Piece,
) -> Vec<PuzzleDense> {
    let mut solutions = Vec::new();
    for rotation in all_rotations(piece) {
        for s in all_puts(already_placed.clone(), piece_count, &rotation) {
            solutions.push(s);
        }
    }
    solutions
}

fn all_rotations(piece: &Piece) -> Vec<Vec<[i32; 3]>> {
    let mut all_rots = Vec::new();
    for theta in 0..4 {
        let mut rotations = Vec::new();
        for [x, y, z] in piece {
            let rotated_z = [
                x * COS[theta] - y * SIN[theta],
                x * SIN[theta] + y * COS[theta],
                *z,
            ];
            rotations.push(rotated_z);
        }
        all_rots.push(rotations);
    }

    let mut all_rotss = all_rots.clone();
    for piece in &all_rots {
        for theta in 0..4 {
            let mut rotations = Vec::new();
            for [x, y, z] in piece {
                let rotated_yz = [
                    x * COS[theta] + z * SIN[theta],
                    *y,
                    -x * SIN[theta] + z * COS[theta],
                ];
                rotations.push(rotated_yz);
            }
            all_rotss.push(rotations);
        }
    }

    let mut all_rotsss = all_rotss.clone();
    for piece in &all_rotss {
        for theta in 0..4 {
            let mut rotations = Vec::new();
            for [x, y, z] in piece {
                let rotated_xyz = [
                    *x,
                    y * COS[theta] - z * SIN[theta],
                    y * SIN[theta] + z * COS[theta],
                ];
                rotations.push(rotated_xyz);
            }
            all_rotsss.push(rotations)
        }
    }

    let mut unique_solutions = HashSet::new();
    let rots = push_to_zero(all_rotsss);
    for mut solution in rots {
        solution.sort();
        unique_solutions.insert(solution);
    }

    let mut rotations = Vec::new();
    for solution in unique_solutions {
        rotations.push(solution);
    }
    rotations
}

fn max_xyz(piece: &Piece) -> (i32, i32, i32) {
    let mut max_x = -1;
    let mut max_y = -1;
    let mut max_z = -1;

    for [x, y, z] in piece {
        if x > &max_x {
            max_x = *x;
        }
        if y > &max_y {
            max_y = *y;
        }
        if z > &max_z {
            max_z = *z;
        }
    }
    (max_x, max_y, max_z)
}

fn all_puts(already_placed: Vec<PuzzleDense>, piece_count: i32, piece: &Piece) -> Vec<PuzzleDense> {
    let mut all_solutions = Vec::new();
    let (max_x, max_y, max_z) = max_xyz(piece);

    for x_step in 0..3 - max_x {
        for y_step in 0..3 - max_y {
            for z_step in 0..3 - max_z {
                let already_placed = already_placed.clone();
                'next_piece: for mut puzzle in already_placed {
                    for [x, y, z] in piece {
                        if puzzle.data[(x + x_step) as usize][(y + y_step) as usize]
                            [(z + z_step) as usize]
                            > 0
                        {
                            continue 'next_piece;
                        }
                    }

                    for [x, y, z] in piece {
                        puzzle.data[(x + x_step) as usize][(y + y_step) as usize]
                            [(z + z_step) as usize] = piece_count;
                    }

                    all_solutions.push(puzzle);
                }
            }
        }
    }

    all_solutions
}

fn zeros() -> [[[i32; 3]; 3]; 3] {
    [
        [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
        [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
        [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
    ]
}

fn push_to_zero(puzzle: Puzzle) -> Puzzle {
    let mut pieces = Vec::new();

    for part in &puzzle {
        let mut piece = Vec::new();
        let mut min_x = 99;
        let mut min_y = 99;
        let mut min_z = 99;

        for [x, y, z] in part {
            if x < &min_x {
                min_x = *x;
            }
            if y < &min_y {
                min_y = *y;
            }
            if z < &min_z {
                min_z = *z;
            }
        }
        for [x, y, z] in part {
            piece.push([x - min_x, y - min_y, z - min_z]);
        }

        pieces.push(piece);
    }

    pieces
}<commit_msg>Make the puzzle argument take a string.<commit_after>extern crate clap;

use clap::Parser;

use std::collections::{HashMap, HashSet};
use std::fmt;
use std::time::Instant;

use cubes_rs::{Piece, Puzzle};

const SIN: [i32; 4] = [0, 1, 0, -1];
const COS: [i32; 4] = [1, 0, -1, 0];

/// Program to display the 3x3 cube solution(s).
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// blue, green, or minotaur.
    #[arg(short, long, default_value_t = String::from("minotaur"))]
    puzzle: String,
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
struct PuzzleDense {
    data: [[[i32; 3]; 3]; 3],
}

impl fmt::Display for PuzzleDense {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "{:?}", self.data[0])?;
        writeln!(f, "{:?}", self.data[1])?;
        writeln!(f, "{:?}", self.data[2])
    }
}

fn main() {
    let now = Instant::now();

    let args = Args::parse();
    let puzzle = match args.puzzle.as_str() {
        "blue" => cubes_rs::blue(),
        "green" => cubes_rs::green(),
        "minotaur" => cubes_rs::minotaur(),
        _ => panic!("invalid argument"),
    };

    let pieces = push_to_zero(puzzle);
    for piece in &pieces {
        for part in piece {
            println!("{:?}", part);
        }
        println!();
    }

    let zero = vec![PuzzleDense { data: zeros() }];

    let mut solutions = Vec::new();
    for (i, piece) in pieces.iter().enumerate() {
        if i == 0 {
            solutions = all_puts(zero.clone(), (i + 1) as i32, piece);
        } else {
            solutions = all_rotations_and_puts(solutions.clone(), (i + 1) as i32, piece);
        }
    }

    let solutions = unique_pieces(solutions);
    for solution in solutions {
        println!("{:}", solution);
    }

    let elapsed_time = now.elapsed();
    println!("Running the program took {} seconds.", elapsed_time.as_millis() as f64 / 1000.0);
}

fn unique_pieces(puzzles: Vec<PuzzleDense>) -> Vec<PuzzleDense> {
    let mut puzzles_unique = HashSet::new();
    for mut puzzle in puzzles {
        let mut puzzle_unique: HashMap<i32, i32> = HashMap::new();
        let mut piece_count = 0;
        for x in 0..3 {
            for y in 0..3 {
                for z in 0..3 {
                    let value = puzzle.data[x][y][z];
                    if puzzle_unique.contains_key(&value) {
                        puzzle.data[x][y][z] = puzzle_unique[&value]
                    } else {
                        piece_count += 1;
                        puzzle_unique.insert(value, piece_count);
                        puzzle.data[x][y][z] = puzzle_unique[&value]
                    }
                }
            }
        }
        puzzles_unique.insert(puzzle);
    }

    let mut all_puzzles = Vec::new();
    for puzzle in puzzles_unique {
        all_puzzles.push(puzzle);
    }
    all_puzzles
}

fn all_rotations_and_puts(
    already_placed: Vec<PuzzleDense>,
    piece_count: i32,
    piece: &Piece,
) -> Vec<PuzzleDense> {
    let mut solutions = Vec::new();
    for rotation in all_rotations(piece) {
        for s in all_puts(already_placed.clone(), piece_count, &rotation) {
            solutions.push(s);
        }
    }
    solutions
}

fn all_rotations(piece: &Piece) -> Vec<Vec<[i32; 3]>> {
    let mut all_rots = Vec::new();
    for theta in 0..4 {
        let mut rotations = Vec::new();
        for [x, y, z] in piece {
            let rotated_z = [
                x * COS[theta] - y * SIN[theta],
                x * SIN[theta] + y * COS[theta],
                *z,
            ];
            rotations.push(rotated_z);
        }
        all_rots.push(rotations);
    }

    let mut all_rotss = all_rots.clone();
    for piece in &all_rots {
        for theta in 0..4 {
            let mut rotations = Vec::new();
            for [x, y, z] in piece {
                let rotated_yz = [
                    x * COS[theta] + z * SIN[theta],
                    *y,
                    -x * SIN[theta] + z * COS[theta],
                ];
                rotations.push(rotated_yz);
            }
            all_rotss.push(rotations);
        }
    }

    let mut all_rotsss = all_rotss.clone();
    for piece in &all_rotss {
        for theta in 0..4 {
            let mut rotations = Vec::new();
            for [x, y, z] in piece {
                let rotated_xyz = [
                    *x,
                    y * COS[theta] - z * SIN[theta],
                    y * SIN[theta] + z * COS[theta],
                ];
                rotations.push(rotated_xyz);
            }
            all_rotsss.push(rotations)
        }
    }

    let mut unique_solutions = HashSet::new();
    let rots = push_to_zero(all_rotsss);
    for mut solution in rots {
        solution.sort();
        unique_solutions.insert(solution);
    }

    let mut rotations = Vec::new();
    for solution in unique_solutions {
        rotations.push(solution);
    }
    rotations
}

fn max_xyz(piece: &Piece) -> (i32, i32, i32) {
    let mut max_x = -1;
    let mut max_y = -1;
    let mut max_z = -1;

    for [x, y, z] in piece {
        if x > &max_x {
            max_x = *x;
        }
        if y > &max_y {
            max_y = *y;
        }
        if z > &max_z {
            max_z = *z;
        }
    }
    (max_x, max_y, max_z)
}

fn all_puts(already_placed: Vec<PuzzleDense>, piece_count: i32, piece: &Piece) -> Vec<PuzzleDense> {
    let mut all_solutions = Vec::new();
    let (max_x, max_y, max_z) = max_xyz(piece);

    for x_step in 0..3 - max_x {
        for y_step in 0..3 - max_y {
            for z_step in 0..3 - max_z {
                let already_placed = already_placed.clone();
                'next_piece: for mut puzzle in already_placed {
                    for [x, y, z] in piece {
                        if puzzle.data[(x + x_step) as usize][(y + y_step) as usize]
                            [(z + z_step) as usize]
                            > 0
                        {
                            continue 'next_piece;
                        }
                    }

                    for [x, y, z] in piece {
                        puzzle.data[(x + x_step) as usize][(y + y_step) as usize]
                            [(z + z_step) as usize] = piece_count;
                    }

                    all_solutions.push(puzzle);
                }
            }
        }
    }

    all_solutions
}

fn zeros() -> [[[i32; 3]; 3]; 3] {
    [
        [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
        [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
        [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
    ]
}

fn push_to_zero(puzzle: Puzzle) -> Puzzle {
    let mut pieces = Vec::new();

    for part in &puzzle {
        let mut piece = Vec::new();
        let mut min_x = 99;
        let mut min_y = 99;
        let mut min_z = 99;

        for [x, y, z] in part {
            if x < &min_x {
                min_x = *x;
            }
            if y < &min_y {
                min_y = *y;
            }
            if z < &min_z {
                min_z = *z;
            }
        }
        for [x, y, z] in part {
            piece.push([x - min_x, y - min_y, z - min_z]);
        }

        pieces.push(piece);
    }

    pieces
}<|endoftext|>