<commit_before>use std::collections::HashMap;
use std::ops::Deref;

use rustc_serialize::{Encodable, Decodable, Encoder, Decoder};

bitflags! {
    pub flags FieldFlags: u32 {
        const FIELD_INDEXED = 0b00000001,
        const FIELD_STORED  = 0b00000010,
    }
}

impl Encodable for FieldFlags {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        let mut flag_strings = Vec::new();

        if self.contains(FIELD_INDEXED) {
            flag_strings.push("INDEXED");
        }

        if self.contains(FIELD_STORED) {
            flag_strings.push("STORED");
        }

        try!(s.emit_str(&flag_strings.join("|")));

        Ok(())
    }
}

impl Decodable for FieldFlags {
    fn decode<D: Decoder>(d: &mut D) -> Result<FieldFlags, D::Error> {
        let s = try!(d.read_str());
        let mut flags = FieldFlags::empty();

        for flag_s in s.split("|") {
            match flag_s {
                "INDEXED" => {
                    flags |= FIELD_INDEXED;
                }
                "STORED" => {
                    flags |= FIELD_STORED;
                }
                _ => {} // TODO: error
            }
        }

        Ok(flags)
    }
}

#[derive(Debug, Clone, PartialEq, RustcEncodable, RustcDecodable)]
pub enum FieldType {
    Text,
    PlainString,
    I64,
    Boolean,
    DateTime,
}

#[derive(Debug, Clone, RustcEncodable, RustcDecodable)]
pub struct FieldInfo {
    name: String,
    pub field_type: FieldType,
    pub field_flags: FieldFlags,
}

impl FieldInfo {
    pub fn new(name: String, field_type: FieldType, field_flags: FieldFlags) -> FieldInfo {
        FieldInfo {
            name: name,
            field_type: field_type,
            field_flags: field_flags,
        }
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
pub struct FieldRef(u32);

impl FieldRef {
    pub fn ord(&self) -> u32 {
        self.0
    }
}

impl Encodable for FieldRef {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        s.emit_u32(self.ord())
    }
}

impl Decodable for FieldRef {
    fn decode<D: Decoder>(d: &mut D) -> Result<FieldRef, D::Error> {
        Ok(FieldRef(try!(d.read_u32())))
    }
}

#[derive(Debug)]
pub enum AddFieldError {
    FieldAlreadyExists(String),
}

#[derive(Debug, Clone, RustcEncodable, RustcDecodable)]
pub struct Schema {
    next_field_id: u32,
    fields: HashMap<FieldRef, FieldInfo>,
    field_names: HashMap<String, FieldRef>,
}

impl Schema {
    pub fn new() -> Schema {
        Schema {
            next_field_id: 1,
            fields: HashMap::new(),
            field_names: HashMap::new(),
        }
    }

    fn new_field_ref(&mut self) -> FieldRef {
        let field_ref = FieldRef(self.next_field_id);
        self.next_field_id += 1;

        field_ref
    }

    pub fn get_field_by_name(&self, name: &str) -> Option<FieldRef> {
        self.field_names.get(name).cloned()
    }

    pub fn add_field(&mut self, name: String, field_type: FieldType, field_flags: FieldFlags) -> Result<FieldRef, AddFieldError> {
        if self.field_names.contains_key(&name) {
            return Err(AddFieldError::FieldAlreadyExists(name));
        }

        let field_ref = self.new_field_ref();
        let field_info = FieldInfo::new(name.clone(), field_type, field_flags);

        self.fields.insert(field_ref, field_info);
        self.field_names.insert(name, field_ref);

        Ok(field_ref)
    }

    pub fn remove_field(&mut self, field_ref: &FieldRef) -> bool {
        match self.fields.remove(field_ref) {
            Some(removed_field) => {
                self.field_names.remove(&removed_field.name);
                true
            }
            None => false
        }
    }
}

impl Deref for Schema {
    type Target = HashMap<FieldRef, FieldInfo>;

    fn deref(&self) -> &HashMap<FieldRef, FieldInfo> {
        &self.fields
    }
}
<commit_msg>Added new method to FieldRef<commit_after>use std::collections::HashMap;
use std::ops::Deref;

use rustc_serialize::{Encodable, Decodable, Encoder, Decoder};

bitflags! {
    pub flags FieldFlags: u32 {
        const FIELD_INDEXED = 0b00000001,
        const FIELD_STORED  = 0b00000010,
    }
}

impl Encodable for FieldFlags {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        let mut flag_strings = Vec::new();

        if self.contains(FIELD_INDEXED) {
            flag_strings.push("INDEXED");
        }

        if self.contains(FIELD_STORED) {
            flag_strings.push("STORED");
        }

        try!(s.emit_str(&flag_strings.join("|")));

        Ok(())
    }
}

impl Decodable for FieldFlags {
    fn decode<D: Decoder>(d: &mut D) -> Result<FieldFlags, D::Error> {
        let s = try!(d.read_str());
        let mut flags = FieldFlags::empty();

        for flag_s in s.split("|") {
            match flag_s {
                "INDEXED" => {
                    flags |= FIELD_INDEXED;
                }
                "STORED" => {
                    flags |= FIELD_STORED;
                }
                _ => {} // TODO: error
            }
        }

        Ok(flags)
    }
}

#[derive(Debug, Clone, PartialEq, RustcEncodable, RustcDecodable)]
pub enum FieldType {
    Text,
    PlainString,
    I64,
    Boolean,
    DateTime,
}

#[derive(Debug, Clone, RustcEncodable, RustcDecodable)]
pub struct FieldInfo {
    name: String,
    pub field_type: FieldType,
    pub field_flags: FieldFlags,
}

impl FieldInfo {
    pub fn new(name: String, field_type: FieldType, field_flags: FieldFlags) -> FieldInfo {
        FieldInfo {
            name: name,
            field_type: field_type,
            field_flags: field_flags,
        }
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
pub struct FieldRef(u32);

impl FieldRef {
    pub fn new(ord: u32) -> FieldRef {
        FieldRef(ord)
    }

    pub fn ord(&self) -> u32 {
        self.0
    }
}

impl Encodable for FieldRef {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        s.emit_u32(self.ord())
    }
}

impl Decodable for FieldRef {
    fn decode<D: Decoder>(d: &mut D) -> Result<FieldRef, D::Error> {
        Ok(FieldRef(try!(d.read_u32())))
    }
}

#[derive(Debug)]
pub enum AddFieldError {
    FieldAlreadyExists(String),
}

#[derive(Debug, Clone, RustcEncodable, RustcDecodable)]
pub struct Schema {
    next_field_id: u32,
    fields: HashMap<FieldRef, FieldInfo>,
    field_names: HashMap<String, FieldRef>,
}

impl Schema {
    pub fn new() -> Schema {
        Schema {
            next_field_id: 1,
            fields: HashMap::new(),
            field_names: HashMap::new(),
        }
    }

    fn new_field_ref(&mut self) -> FieldRef {
        let field_ref = FieldRef(self.next_field_id);
        self.next_field_id += 1;

        field_ref
    }

    pub fn get_field_by_name(&self, name: &str) -> Option<FieldRef> {
        self.field_names.get(name).cloned()
    }

    pub fn add_field(&mut self, name: String, field_type: FieldType, field_flags: FieldFlags) -> Result<FieldRef, AddFieldError> {
        if self.field_names.contains_key(&name) {
            return Err(AddFieldError::FieldAlreadyExists(name));
        }

        let field_ref = self.new_field_ref();
        let field_info = FieldInfo::new(name.clone(), field_type, field_flags);

        self.fields.insert(field_ref, field_info);
        self.field_names.insert(name, field_ref);

        Ok(field_ref)
    }

    pub fn remove_field(&mut self, field_ref: &FieldRef) -> bool {
        match self.fields.remove(field_ref) {
            Some(removed_field) => {
                self.field_names.remove(&removed_field.name);
                true
            }
            None => false
        }
    }
}

impl Deref for Schema {
    type Target = HashMap<FieldRef, FieldInfo>;

    fn deref(&self) -> &HashMap<FieldRef, FieldInfo> {
        &self.fields
    }
}
<|endoftext|>


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

use crate::common::types::{
    CliError, CliTypedResult, MovePackageDir, PoolAddressArgs, PromptOptions, TransactionOptions,
    TransactionSummary,
};
use crate::common::utils::prompt_yes_with_override;
#[cfg(feature = "no-upload-proposal")]
use crate::common::utils::read_from_file;
use crate::move_tool::{FrameworkPackageArgs, IncludedArtifacts};
use crate::{CliCommand, CliResult};
use aptos_crypto::HashValue;
use aptos_logger::warn;
use aptos_rest_client::aptos_api_types::U64;
use aptos_rest_client::Transaction;
use aptos_types::{
    account_address::AccountAddress,
    transaction::{Script, TransactionPayload},
};
use async_trait::async_trait;
use cached_packages::aptos_stdlib;
use clap::Parser;
use framework::{BuildOptions, BuiltPackage, ReleasePackage};
use move_deps::move_core_types::transaction_argument::TransactionArgument;
use reqwest::Url;
use serde::Deserialize;
use serde::Serialize;
use std::path::Path;
use std::{collections::BTreeMap, fmt::Formatter, fs, path::PathBuf};
use tempfile::TempDir;

/// Tool for on-chain governance
///
/// This tool allows voters that have stake to vote the ability to
/// propose changes to the chain, as well as vote and execute these
/// proposals.
#[derive(Parser)]
pub enum GovernanceTool {
    Propose(SubmitProposal),
    Vote(SubmitVote),
    ExecuteProposal(ExecuteProposal),
    GenerateUpgradeProposal(GenerateUpgradeProposal),
}

impl GovernanceTool {
    pub async fn execute(self) -> CliResult {
        use GovernanceTool::*;
        match self {
            Propose(tool) => tool.execute_serialized().await,
            Vote(tool) => tool.execute_serialized().await,
            ExecuteProposal(tool) => tool.execute_serialized().await,
            GenerateUpgradeProposal(tool) => tool.execute_serialized_success().await,
        }
    }
}

/// Submit proposal to other validators to be proposed on
#[derive(Parser)]
pub struct SubmitProposal {
    /// Location of the JSON metadata of the proposal
    #[clap(long)]
    pub(crate) metadata_url: Url,

    #[cfg(feature = "no-upload-proposal")]
    /// A JSON file to be uploaded later at the metadata URL
    #[clap(long)]
    pub(crate) metadata_path: Option<PathBuf>,

    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
    #[clap(flatten)]
    pub(crate) pool_address_args: PoolAddressArgs,
    #[clap(flatten)]
    pub(crate) compile_proposal_args: CompileScriptFunction,
}

#[async_trait]
impl CliCommand<ProposalSubmissionSummary> for SubmitProposal {
    fn command_name(&self) -> &'static str {
        "SubmitProposal"
    }

    async fn execute(mut self) -> CliTypedResult<ProposalSubmissionSummary> {
        let (_bytecode, script_hash) = self
            .compile_proposal_args
            .compile("SubmitProposal", self.txn_options.prompt_options)?;

        // Validate the proposal metadata
        let (metadata, metadata_hash) = self.get_metadata().await?;

        println!(
            "{}\n\tMetadata Hash: {}\n\tScript Hash: {}",
            metadata, metadata_hash, script_hash
        );
        prompt_yes_with_override(
            "Do you want to submit this proposal?",
            self.txn_options.prompt_options,
        )?;

        let txn = self
            .txn_options
            .submit_transaction(aptos_stdlib::aptos_governance_create_proposal(
                self.pool_address_args.pool_address,
                script_hash.to_vec(),
                self.metadata_url.to_string().as_bytes().to_vec(),
                metadata_hash.to_hex().as_bytes().to_vec(),
            ))
            .await?;
        let txn_summary = TransactionSummary::from(&txn);
        if let Transaction::UserTransaction(inner) = txn {
            // Find event with proposal id
            let proposal_id = if let Some(event) = inner.events.into_iter().find(|event| {
                event.typ.to_string().as_str() == "0x1::aptos_governance::CreateProposalEvent"
            }) {
                let data: CreateProposalEvent =
                    serde_json::from_value(event.data).map_err(|_| {
                        CliError::UnexpectedError(
                            "Failed to parse Proposal event to get ProposalId".to_string(),
                        )
                    })?;
                Some(data.proposal_id.0)
            } else {
                warn!("No proposal event found to find proposal id");
                None
            };

            return Ok(ProposalSubmissionSummary {
                proposal_id,
                transaction: txn_summary,
            });
        }
        Err(CliError::UnexpectedError(
            "Unable to find parse proposal transaction output".to_string(),
        ))
    }
}

impl SubmitProposal {
    /// Retrieve metadata and validate it
    async fn get_metadata(&self) -> CliTypedResult<(ProposalMetadata, HashValue)> {
        #[cfg(feature = "no-upload-proposal")]
        let bytes = if let Some(ref path) = self.metadata_path {
            read_from_file(path)?
        } else {
            get_metadata_from_url(&self.metadata_url).await?
        };
        #[cfg(not(feature = "no-upload-proposal"))]
        let bytes = get_metadata_from_url(&self.metadata_url).await?;

        let metadata: ProposalMetadata = serde_json::from_slice(&bytes).map_err(|err| {
            CliError::CommandArgumentError(format!(
                "Metadata is not in a proper JSON format: {}",
                err
            ))
        })?;
        Url::parse(&metadata.source_code_url).map_err(|err| {
            CliError::CommandArgumentError(format!(
                "Source code URL {} is invalid {}",
                metadata.source_code_url, err
            ))
        })?;
        Url::parse(&metadata.discussion_url).map_err(|err| {
            CliError::CommandArgumentError(format!(
                "Discussion URL {} is invalid {}",
                metadata.discussion_url, err
            ))
        })?;
        let metadata_hash = HashValue::sha3_256_of(&bytes);
        Ok((metadata, metadata_hash))
    }
}

async fn get_metadata_from_url(metadata_url: &Url) -> CliTypedResult<Vec<u8>> {
    let client = reqwest::ClientBuilder::default()
        .tls_built_in_root_certs(true)
        .build()
        .map_err(|err| CliError::UnexpectedError(format!("Failed to build HTTP client {}", err)))?;
    client
        .get(metadata_url.clone())
        .send()
        .await
        .map_err(|err| {
            CliError::CommandArgumentError(format!(
                "Failed to fetch metadata url {}: {}",
                metadata_url, err
            ))
        })?
        .bytes()
        .await
        .map(|b| b.to_vec())
        .map_err(|err| {
            CliError::CommandArgumentError(format!(
                "Failed to fetch metadata url {}: {}",
                metadata_url, err
            ))
        })
}

#[derive(Debug, Deserialize, Serialize)]
struct CreateProposalEvent {
    proposal_id: U64,
}

#[derive(Debug, Deserialize, Serialize)]
struct ProposalSubmissionSummary {
    proposal_id: Option<u64>,
    #[serde(flatten)]
    transaction: TransactionSummary,
}

/// Submit a vote on a current proposal
#[derive(Parser)]
pub struct SubmitVote {
    /// Id of proposal to vote on
    #[clap(long)]
    pub(crate) proposal_id: u64,

    /// Vote yes on the proposal
    #[clap(long, group = "vote")]
    pub(crate) yes: bool,

    /// Vote no on the proposal
    #[clap(long, group = "vote")]
    pub(crate) no: bool,

    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
    #[clap(flatten)]
    pub(crate) pool_address_args: PoolAddressArgs,
}

#[async_trait]
impl CliCommand<TransactionSummary> for SubmitVote {
    fn command_name(&self) -> &'static str {
        "SubmitVote"
    }

    async fn execute(mut self) -> CliTypedResult<TransactionSummary> {
        let (vote_str, vote) = match (self.yes, self.no) {
            (true, false) => ("Yes", true),
            (false, true) => ("No", false),
            (_, _) => {
                return Err(CliError::CommandArgumentError(
                    "Must choose either --yes or --no".to_string(),
                ))
            }
        };

        // TODO: Display details of proposal

        prompt_yes_with_override(
            &format!("Are you sure you want to vote {}", vote_str),
            self.txn_options.prompt_options,
        )?;

        self.txn_options
            .submit_transaction(aptos_stdlib::aptos_governance_vote(
                self.pool_address_args.pool_address,
                self.proposal_id,
                vote,
            ))
            .await
            .map(TransactionSummary::from)
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ProposalMetadata {
    title: String,
    description: String,
    source_code_url: String,
    discussion_url: String,
}

impl std::fmt::Display for ProposalMetadata {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Proposal:\n\tTitle:{}\n\tDescription:{}\n\tSource code URL:{}\n\tDiscussion URL:{}",
            self.title, self.description, self.source_code_url, self.discussion_url
        )
    }
}

fn compile_in_temp_dir(
    script_name: &str,
    script_path: &Path,
    framework_package_args: &FrameworkPackageArgs,
    prompt_options: PromptOptions,
) -> CliTypedResult<(Vec<u8>, HashValue)> {
    // Make a temporary directory for compilation
    let temp_dir = TempDir::new().map_err(|err| {
        CliError::UnexpectedError(format!("Failed to create temporary directory {}", err))
    })?;

    // Initialize a move directory
    let package_dir = temp_dir.path();
    framework_package_args.init_move_dir(
        package_dir,
        script_name,
        BTreeMap::new(),
        prompt_options,
    )?;

    // Insert the new script
    let sources_dir = package_dir.join("sources");
    let new_script_path = if let Some(file_name) = script_path.file_name() {
        sources_dir.join(file_name)
    } else {
        // If for some reason we can't get the move file
        sources_dir.join("script.move")
    };
    fs::copy(script_path, new_script_path.as_path()).map_err(|err| {
        CliError::IO(
            format!(
                "Failed to copy {} to {}",
                script_path.display(),
                new_script_path.display()
            ),
            err,
        )
    })?;

    // Compile the script
    compile_script(package_dir)
}

fn compile_script(package_dir: &Path) -> CliTypedResult<(Vec<u8>, HashValue)> {
    let build_options = BuildOptions {
        with_srcs: false,
        with_abis: false,
        with_source_maps: false,
        with_error_map: false,
        install_dir: None,
        named_addresses: Default::default(),
    };

    let pack = BuiltPackage::build(package_dir.to_path_buf(), build_options)?;

    let scripts_count = pack.script_count();

    if scripts_count != 1 {
        return Err(CliError::UnexpectedError(format!(
            "Only one script can be prepared a time. Make sure one and only one script file \
                is included in the Move package. Found {} scripts.",
            scripts_count
        )));
    }

    let bytes = pack.extract_script_code().pop().unwrap();
    let hash = HashValue::sha3_256_of(bytes.as_slice());
    Ok((bytes, hash))
}

/// Execute a proposal that has passed voting requirements
#[derive(Parser)]
pub struct ExecuteProposal {
    /// Proposal Id being executed
    #[clap(long)]
    pub(crate) proposal_id: u64,

    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
    #[clap(flatten)]
    pub(crate) compile_proposal_args: CompileScriptFunction,
}

#[async_trait]
impl CliCommand<TransactionSummary> for ExecuteProposal {
    fn command_name(&self) -> &'static str {
        "ExecuteProposal"
    }

    async fn execute(mut self) -> CliTypedResult<TransactionSummary> {
        let (bytecode, _script_hash) = self
            .compile_proposal_args
            .compile("ExecuteProposal", self.txn_options.prompt_options)?;
        // TODO: Check hash so we don't do a failed roundtrip?

        let args = vec![TransactionArgument::U64(self.proposal_id)];
        let txn = TransactionPayload::Script(Script::new(bytecode, vec![], args));

        self.txn_options
            .submit_transaction(txn)
            .await
            .map(TransactionSummary::from)
    }
}

/// Execute a proposal that has passed voting requirements
#[derive(Parser)]
pub struct CompileScriptFunction {
    /// Path to the Move script for the proposal
    #[clap(long, parse(from_os_str))]
    pub script_path: PathBuf,

    #[clap(flatten)]
    pub(crate) framework_package_args: FrameworkPackageArgs,
}

impl CompileScriptFunction {
    pub(crate) fn compile(
        &self,
        script_name: &str,
        prompt_options: PromptOptions,
    ) -> CliTypedResult<(Vec<u8>, HashValue)> {
        // Check script file
        let script_path = self.script_path.as_path();
        if !self.script_path.exists() {
            return Err(CliError::CommandArgumentError(format!(
                "{} does not exist",
                script_path.display()
            )));
        } else if self.script_path.is_dir() {
            return Err(CliError::CommandArgumentError(format!(
                "{} is a directory",
                script_path.display()
            )));
        }

        // Compile script
        compile_in_temp_dir(
            script_name,
            script_path,
            &self.framework_package_args,
            prompt_options,
        )
    }
}

/// Generates a package upgrade proposal script.
#[derive(Parser)]
pub struct GenerateUpgradeProposal {
    /// Address of the account which the proposal addresses.
    #[clap(long, parse(try_from_str=crate::common::types::load_account_arg))]
    pub(crate) account: AccountAddress,

    /// Where to store the generated proposal
    #[clap(long, parse(from_os_str), default_value = "proposal.move")]
    pub(crate) output: PathBuf,

    /// What artifacts to include in the package. This can be one of `none`, `sparse`, and
    /// `all`. `none` is the most compact form and does not allow to reconstruct a source
    /// package from chain; `sparse` is the minimal set of artifacts needed to reconstruct
    /// a source package; `all` includes all available artifacts. The choice of included
    /// artifacts heavily influences the size and therefore gas cost of publishing: `none`
    /// is the size of bytecode alone; `sparse` is roughly 2 times as much; and `all` 3-4
    /// as much.
    #[clap(long, default_value_t = IncludedArtifacts::Sparse)]
    pub(crate) included_artifacts: IncludedArtifacts,

    #[clap(flatten)]
    pub(crate) move_options: MovePackageDir,
}

#[async_trait]
impl CliCommand<()> for GenerateUpgradeProposal {
    fn command_name(&self) -> &'static str {
        "GenerateUpgradeProposal"
    }

    async fn execute(self) -> CliTypedResult<()> {
        let GenerateUpgradeProposal {
            move_options,
            account,
            included_artifacts,
            output,
        } = self;
        let package_path = move_options.get_package_path()?;
        let options = included_artifacts.build_options(move_options.named_addresses());
        let package = BuiltPackage::build(package_path, options)?;
        let release = ReleasePackage::new(package)?;
        release.generate_script_proposal(account, output)?;
        Ok(())
    }
}
<commit_msg>[cli] add support for (already) compiled scripts<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::common::types::{
    CliError, CliTypedResult, MovePackageDir, PoolAddressArgs, PromptOptions, TransactionOptions,
    TransactionSummary,
};
use crate::common::utils::prompt_yes_with_override;
#[cfg(feature = "no-upload-proposal")]
use crate::common::utils::read_from_file;
use crate::move_tool::{FrameworkPackageArgs, IncludedArtifacts};
use crate::{CliCommand, CliResult};
use aptos_crypto::HashValue;
use aptos_logger::warn;
use aptos_rest_client::aptos_api_types::U64;
use aptos_rest_client::Transaction;
use aptos_types::{
    account_address::AccountAddress,
    transaction::{Script, TransactionPayload},
};
use async_trait::async_trait;
use cached_packages::aptos_stdlib;
use clap::Parser;
use framework::{BuildOptions, BuiltPackage, ReleasePackage};
use move_deps::move_core_types::transaction_argument::TransactionArgument;
use reqwest::Url;
use serde::Deserialize;
use serde::Serialize;
use std::path::Path;
use std::{collections::BTreeMap, fmt::Formatter, fs, path::PathBuf};
use tempfile::TempDir;

/// Tool for on-chain governance
///
/// This tool allows voters that have stake to vote the ability to
/// propose changes to the chain, as well as vote and execute these
/// proposals.
#[derive(Parser)]
pub enum GovernanceTool {
    Propose(SubmitProposal),
    Vote(SubmitVote),
    ExecuteProposal(ExecuteProposal),
    GenerateUpgradeProposal(GenerateUpgradeProposal),
}

impl GovernanceTool {
    pub async fn execute(self) -> CliResult {
        use GovernanceTool::*;
        match self {
            Propose(tool) => tool.execute_serialized().await,
            Vote(tool) => tool.execute_serialized().await,
            ExecuteProposal(tool) => tool.execute_serialized().await,
            GenerateUpgradeProposal(tool) => tool.execute_serialized_success().await,
        }
    }
}

/// Submit proposal to other validators to be proposed on
#[derive(Parser)]
pub struct SubmitProposal {
    /// Location of the JSON metadata of the proposal
    #[clap(long)]
    pub(crate) metadata_url: Url,

    #[cfg(feature = "no-upload-proposal")]
    /// A JSON file to be uploaded later at the metadata URL
    #[clap(long)]
    pub(crate) metadata_path: Option<PathBuf>,

    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
    #[clap(flatten)]
    pub(crate) pool_address_args: PoolAddressArgs,
    #[clap(flatten)]
    pub(crate) compile_proposal_args: CompileScriptFunction,
}

#[async_trait]
impl CliCommand<ProposalSubmissionSummary> for SubmitProposal {
    fn command_name(&self) -> &'static str {
        "SubmitProposal"
    }

    async fn execute(mut self) -> CliTypedResult<ProposalSubmissionSummary> {
        let (_bytecode, script_hash) = self
            .compile_proposal_args
            .compile("SubmitProposal", self.txn_options.prompt_options)?;

        // Validate the proposal metadata
        let (metadata, metadata_hash) = self.get_metadata().await?;

        println!(
            "{}\n\tMetadata Hash: {}\n\tScript Hash: {}",
            metadata, metadata_hash, script_hash
        );
        prompt_yes_with_override(
            "Do you want to submit this proposal?",
            self.txn_options.prompt_options,
        )?;

        let txn = self
            .txn_options
            .submit_transaction(aptos_stdlib::aptos_governance_create_proposal(
                self.pool_address_args.pool_address,
                script_hash.to_vec(),
                self.metadata_url.to_string().as_bytes().to_vec(),
                metadata_hash.to_hex().as_bytes().to_vec(),
            ))
            .await?;
        let txn_summary = TransactionSummary::from(&txn);
        if let Transaction::UserTransaction(inner) = txn {
            // Find event with proposal id
            let proposal_id = if let Some(event) = inner.events.into_iter().find(|event| {
                event.typ.to_string().as_str() == "0x1::aptos_governance::CreateProposalEvent"
            }) {
                let data: CreateProposalEvent =
                    serde_json::from_value(event.data).map_err(|_| {
                        CliError::UnexpectedError(
                            "Failed to parse Proposal event to get ProposalId".to_string(),
                        )
                    })?;
                Some(data.proposal_id.0)
            } else {
                warn!("No proposal event found to find proposal id");
                None
            };

            return Ok(ProposalSubmissionSummary {
                proposal_id,
                transaction: txn_summary,
            });
        }
        Err(CliError::UnexpectedError(
            "Unable to find parse proposal transaction output".to_string(),
        ))
    }
}

impl SubmitProposal {
    /// Retrieve metadata and validate it
    async fn get_metadata(&self) -> CliTypedResult<(ProposalMetadata, HashValue)> {
        #[cfg(feature = "no-upload-proposal")]
        let bytes = if let Some(ref path) = self.metadata_path {
            read_from_file(path)?
        } else {
            get_metadata_from_url(&self.metadata_url).await?
        };
        #[cfg(not(feature = "no-upload-proposal"))]
        let bytes = get_metadata_from_url(&self.metadata_url).await?;

        let metadata: ProposalMetadata = serde_json::from_slice(&bytes).map_err(|err| {
            CliError::CommandArgumentError(format!(
                "Metadata is not in a proper JSON format: {}",
                err
            ))
        })?;
        Url::parse(&metadata.source_code_url).map_err(|err| {
            CliError::CommandArgumentError(format!(
                "Source code URL {} is invalid {}",
                metadata.source_code_url, err
            ))
        })?;
        Url::parse(&metadata.discussion_url).map_err(|err| {
            CliError::CommandArgumentError(format!(
                "Discussion URL {} is invalid {}",
                metadata.discussion_url, err
            ))
        })?;
        let metadata_hash = HashValue::sha3_256_of(&bytes);
        Ok((metadata, metadata_hash))
    }
}

async fn get_metadata_from_url(metadata_url: &Url) -> CliTypedResult<Vec<u8>> {
    let client = reqwest::ClientBuilder::default()
        .tls_built_in_root_certs(true)
        .build()
        .map_err(|err| CliError::UnexpectedError(format!("Failed to build HTTP client {}", err)))?;
    client
        .get(metadata_url.clone())
        .send()
        .await
        .map_err(|err| {
            CliError::CommandArgumentError(format!(
                "Failed to fetch metadata url {}: {}",
                metadata_url, err
            ))
        })?
        .bytes()
        .await
        .map(|b| b.to_vec())
        .map_err(|err| {
            CliError::CommandArgumentError(format!(
                "Failed to fetch metadata url {}: {}",
                metadata_url, err
            ))
        })
}

#[derive(Debug, Deserialize, Serialize)]
struct CreateProposalEvent {
    proposal_id: U64,
}

#[derive(Debug, Deserialize, Serialize)]
struct ProposalSubmissionSummary {
    proposal_id: Option<u64>,
    #[serde(flatten)]
    transaction: TransactionSummary,
}

/// Submit a vote on a current proposal
#[derive(Parser)]
pub struct SubmitVote {
    /// Id of proposal to vote on
    #[clap(long)]
    pub(crate) proposal_id: u64,

    /// Vote yes on the proposal
    #[clap(long, group = "vote")]
    pub(crate) yes: bool,

    /// Vote no on the proposal
    #[clap(long, group = "vote")]
    pub(crate) no: bool,

    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
    #[clap(flatten)]
    pub(crate) pool_address_args: PoolAddressArgs,
}

#[async_trait]
impl CliCommand<TransactionSummary> for SubmitVote {
    fn command_name(&self) -> &'static str {
        "SubmitVote"
    }

    async fn execute(mut self) -> CliTypedResult<TransactionSummary> {
        let (vote_str, vote) = match (self.yes, self.no) {
            (true, false) => ("Yes", true),
            (false, true) => ("No", false),
            (_, _) => {
                return Err(CliError::CommandArgumentError(
                    "Must choose either --yes or --no".to_string(),
                ))
            }
        };

        // TODO: Display details of proposal

        prompt_yes_with_override(
            &format!("Are you sure you want to vote {}", vote_str),
            self.txn_options.prompt_options,
        )?;

        self.txn_options
            .submit_transaction(aptos_stdlib::aptos_governance_vote(
                self.pool_address_args.pool_address,
                self.proposal_id,
                vote,
            ))
            .await
            .map(TransactionSummary::from)
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ProposalMetadata {
    title: String,
    description: String,
    source_code_url: String,
    discussion_url: String,
}

impl std::fmt::Display for ProposalMetadata {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Proposal:\n\tTitle:{}\n\tDescription:{}\n\tSource code URL:{}\n\tDiscussion URL:{}",
            self.title, self.description, self.source_code_url, self.discussion_url
        )
    }
}

fn compile_in_temp_dir(
    script_name: &str,
    script_path: &Path,
    framework_package_args: &FrameworkPackageArgs,
    prompt_options: PromptOptions,
) -> CliTypedResult<(Vec<u8>, HashValue)> {
    // Make a temporary directory for compilation
    let temp_dir = TempDir::new().map_err(|err| {
        CliError::UnexpectedError(format!("Failed to create temporary directory {}", err))
    })?;

    // Initialize a move directory
    let package_dir = temp_dir.path();
    framework_package_args.init_move_dir(
        package_dir,
        script_name,
        BTreeMap::new(),
        prompt_options,
    )?;

    // Insert the new script
    let sources_dir = package_dir.join("sources");
    let new_script_path = if let Some(file_name) = script_path.file_name() {
        sources_dir.join(file_name)
    } else {
        // If for some reason we can't get the move file
        sources_dir.join("script.move")
    };
    fs::copy(script_path, new_script_path.as_path()).map_err(|err| {
        CliError::IO(
            format!(
                "Failed to copy {} to {}",
                script_path.display(),
                new_script_path.display()
            ),
            err,
        )
    })?;

    // Compile the script
    compile_script(package_dir)
}

fn compile_script(package_dir: &Path) -> CliTypedResult<(Vec<u8>, HashValue)> {
    let build_options = BuildOptions {
        with_srcs: false,
        with_abis: false,
        with_source_maps: false,
        with_error_map: false,
        install_dir: None,
        named_addresses: Default::default(),
    };

    let pack = BuiltPackage::build(package_dir.to_path_buf(), build_options)?;

    let scripts_count = pack.script_count();

    if scripts_count != 1 {
        return Err(CliError::UnexpectedError(format!(
            "Only one script can be prepared a time. Make sure one and only one script file \
                is included in the Move package. Found {} scripts.",
            scripts_count
        )));
    }

    let bytes = pack.extract_script_code().pop().unwrap();
    let hash = HashValue::sha3_256_of(bytes.as_slice());
    Ok((bytes, hash))
}

/// Execute a proposal that has passed voting requirements
#[derive(Parser)]
pub struct ExecuteProposal {
    /// Proposal Id being executed
    #[clap(long)]
    pub(crate) proposal_id: u64,

    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
    #[clap(flatten)]
    pub(crate) compile_proposal_args: CompileScriptFunction,
}

#[async_trait]
impl CliCommand<TransactionSummary> for ExecuteProposal {
    fn command_name(&self) -> &'static str {
        "ExecuteProposal"
    }

    async fn execute(mut self) -> CliTypedResult<TransactionSummary> {
        let (bytecode, _script_hash) = self
            .compile_proposal_args
            .compile("ExecuteProposal", self.txn_options.prompt_options)?;
        // TODO: Check hash so we don't do a failed roundtrip?

        let args = vec![TransactionArgument::U64(self.proposal_id)];
        let txn = TransactionPayload::Script(Script::new(bytecode, vec![], args));

        self.txn_options
            .submit_transaction(txn)
            .await
            .map(TransactionSummary::from)
    }
}

/// Execute a proposal that has passed voting requirements
#[derive(Parser)]
pub struct CompileScriptFunction {
    /// Path to the Move script for the proposal
    #[clap(long, group = "script", parse(from_os_str))]
    pub script_path: Option<PathBuf>,

    /// Path to the Move script for the proposal
    #[clap(long, group = "script", parse(from_os_str))]
    pub compiled_script_path: Option<PathBuf>,

    #[clap(flatten)]
    pub(crate) framework_package_args: FrameworkPackageArgs,
}

impl CompileScriptFunction {
    pub(crate) fn compile(
        &self,
        script_name: &str,
        prompt_options: PromptOptions,
    ) -> CliTypedResult<(Vec<u8>, HashValue)> {
        if let Some(compiled_script_path) = &self.compiled_script_path {
            let bytes = std::fs::read(compiled_script_path).map_err(|e| {
                CliError::IO(format!("Unable to read {:?}", self.compiled_script_path), e)
            })?;
            let hash = HashValue::sha3_256_of(bytes.as_slice());
            return Ok((bytes, hash));
        }

        // Check script file
        let script_path = self
            .script_path
            .as_ref()
            .ok_or_else(|| {
                CliError::CommandArgumentError(
                    "Must choose either --compiled-script-path or --script-path".to_string(),
                )
            })?
            .as_path();
        if !script_path.exists() {
            return Err(CliError::CommandArgumentError(format!(
                "{} does not exist",
                script_path.display()
            )));
        } else if script_path.is_dir() {
            return Err(CliError::CommandArgumentError(format!(
                "{} is a directory",
                script_path.display()
            )));
        }

        // Compile script
        compile_in_temp_dir(
            script_name,
            script_path,
            &self.framework_package_args,
            prompt_options,
        )
    }
}

/// Generates a package upgrade proposal script.
#[derive(Parser)]
pub struct GenerateUpgradeProposal {
    /// Address of the account which the proposal addresses.
    #[clap(long, parse(try_from_str=crate::common::types::load_account_arg))]
    pub(crate) account: AccountAddress,

    /// Where to store the generated proposal
    #[clap(long, parse(from_os_str), default_value = "proposal.move")]
    pub(crate) output: PathBuf,

    /// What artifacts to include in the package. This can be one of `none`, `sparse`, and
    /// `all`. `none` is the most compact form and does not allow to reconstruct a source
    /// package from chain; `sparse` is the minimal set of artifacts needed to reconstruct
    /// a source package; `all` includes all available artifacts. The choice of included
    /// artifacts heavily influences the size and therefore gas cost of publishing: `none`
    /// is the size of bytecode alone; `sparse` is roughly 2 times as much; and `all` 3-4
    /// as much.
    #[clap(long, default_value_t = IncludedArtifacts::Sparse)]
    pub(crate) included_artifacts: IncludedArtifacts,

    #[clap(flatten)]
    pub(crate) move_options: MovePackageDir,
}

#[async_trait]
impl CliCommand<()> for GenerateUpgradeProposal {
    fn command_name(&self) -> &'static str {
        "GenerateUpgradeProposal"
    }

    async fn execute(self) -> CliTypedResult<()> {
        let GenerateUpgradeProposal {
            move_options,
            account,
            included_artifacts,
            output,
        } = self;
        let package_path = move_options.get_package_path()?;
        let options = included_artifacts.build_options(move_options.named_addresses());
        let package = BuiltPackage::build(package_path, options)?;
        let release = ReleasePackage::new(package)?;
        release.generate_script_proposal(account, output)?;
        Ok(())
    }
}
<|endoftext|>


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

use crate::common::types::{
    CliError, CliTypedResult, PoolAddressArgs, PromptOptions, TransactionOptions,
    TransactionSummary,
};
use crate::common::utils::prompt_yes_with_override;
use crate::move_tool::{compile_move, init_move_dir, ArgWithType, FunctionArgType};
use crate::{CliCommand, CliResult};
use aptos_crypto::HashValue;
use aptos_rest_client::{aptos_api_types::MoveType, Transaction};
use aptos_types::{
    account_address::AccountAddress,
    transaction::{Script, TransactionPayload},
};
use async_trait::async_trait;
use clap::Parser;
use move_deps::{
    move_compiler::compiled_unit::CompiledUnitEnum,
    move_core_types::{language_storage::TypeTag, transaction_argument::TransactionArgument},
    move_package::BuildConfig,
};
use reqwest::Url;
use serde::Deserialize;
use serde::Serialize;
use std::path::Path;
use std::{
    collections::BTreeMap,
    convert::{TryFrom, TryInto},
    fmt::Formatter,
    fs,
    path::PathBuf,
};
use tempfile::TempDir;

/// Tool for on-chain governance
///
/// This tool allows voters that have stake to vote the ability to
/// propose changes to the chain, as well as vote and execute these
/// proposals.
#[derive(Parser)]
pub enum GovernanceTool {
    Propose(SubmitProposal),
    Vote(SubmitVote),
    ExecuteProposal(ExecuteProposal),
}

impl GovernanceTool {
    pub async fn execute(self) -> CliResult {
        use GovernanceTool::*;
        match self {
            Propose(tool) => tool.execute_serialized().await,
            Vote(tool) => tool.execute_serialized().await,
            ExecuteProposal(tool) => tool.execute_serialized().await,
        }
    }
}

/// Submit proposal to other validators to be proposed on
#[derive(Parser)]
pub struct SubmitProposal {
    /// Path to the Move script for the proposal
    #[clap(long, parse(from_os_str))]
    pub script_path: PathBuf,

    /// Git hash or branch of the framework in aptos core
    #[clap(long)]
    pub framework_git_rev: String,

    /// Code location of the script to be voted on
    #[clap(long)]
    pub(crate) metadata_url: Url,

    #[clap(flatten)]
    pub(crate) prompt_options: PromptOptions,
    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
    #[clap(flatten)]
    pub(crate) pool_address_args: PoolAddressArgs,
}

#[async_trait]
impl CliCommand<Transaction> for SubmitProposal {
    fn command_name(&self) -> &'static str {
        "SubmitProposal"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        // Check script file
        let script_path = self.script_path.as_path();
        if self.script_path.exists() {
            return Err(CliError::CommandArgumentError(format!(
                "{} does not exist",
                script_path.display()
            )));
        } else if self.script_path.is_dir() {
            return Err(CliError::CommandArgumentError(format!(
                "{} is a directory",
                script_path.display()
            )));
        }

        // Compile script
        let (_bytecode, script_hash) =
            compile_in_temp_dir(script_path, &self.framework_git_rev, self.prompt_options)?;

        // Validate the proposal metadata
        let (metadata, metadata_hash) = get_metadata(self.metadata_url.clone()).await?;

        println!(
            "{}\n\tMetadata Hash: {}\n\tScript Hash: {}",
            metadata, metadata_hash, script_hash
        );
        prompt_yes_with_override("Do you want to submit this proposal?", self.prompt_options)?;

        self.txn_options
            .submit_script_function(
                AccountAddress::ONE,
                "aptos_governance",
                "create_proposal",
                vec![],
                vec![
                    bcs::to_bytes(&self.pool_address_args.pool_address)?,
                    bcs::to_bytes(&script_hash)?,
                    bcs::to_bytes(&self.metadata_url.to_string())?,
                    bcs::to_bytes(&metadata_hash.to_hex())?,
                ],
            )
            .await
    }
}

/// Retrieve metadata and validate it
async fn get_metadata(metadata_url: Url) -> CliTypedResult<(ProposalMetadata, HashValue)> {
    let client = reqwest::ClientBuilder::default()
        .tls_built_in_root_certs(true)
        .build()
        .map_err(|err| CliError::UnexpectedError(format!("Failed to build HTTP client {}", err)))?;
    let bytes = client
        .get(metadata_url.clone())
        .send()
        .await
        .map_err(|err| {
            CliError::CommandArgumentError(format!(
                "Failed to fetch metadata url {}: {}",
                metadata_url, err
            ))
        })?
        .bytes()
        .await
        .map_err(|err| {
            CliError::CommandArgumentError(format!(
                "Failed to fetch metadata url {}: {}",
                metadata_url, err
            ))
        })?;
    let metadata: ProposalMetadata = serde_json::from_slice(&bytes).map_err(|err| {
        CliError::CommandArgumentError(format!("Metadata is not in a proper JSON format: {}", err))
    })?;
    Url::parse(&metadata.source_code_url).map_err(|err| {
        CliError::CommandArgumentError(format!(
            "Source code URL {} is invalid {}",
            metadata.source_code_url, err
        ))
    })?;
    Url::parse(&metadata.discussion_url).map_err(|err| {
        CliError::CommandArgumentError(format!(
            "Discussion URL {} is invalid {}",
            metadata.discussion_url, err
        ))
    })?;
    let metadata_hash = HashValue::sha3_256_of(&bytes);
    Ok((metadata, metadata_hash))
}

/// Submit a vote on a current proposal
#[derive(Parser)]
pub struct SubmitVote {
    /// Id of proposal to vote on
    #[clap(long)]
    pub(crate) proposal_id: u64,

    /// Vote yes on the proposal
    #[clap(long, group = "vote")]
    pub(crate) yes: bool,

    /// Vote no on the proposal
    #[clap(long, group = "vote")]
    pub(crate) no: bool,

    #[clap(flatten)]
    pub(crate) prompt_options: PromptOptions,
    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
    #[clap(flatten)]
    pub(crate) pool_address_args: PoolAddressArgs,
}

#[async_trait]
impl CliCommand<Transaction> for SubmitVote {
    fn command_name(&self) -> &'static str {
        "SubmitVote"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        let (vote_str, vote) = match (self.yes, self.no) {
            (true, false) => ("Yes", true),
            (false, true) => ("No", false),
            (_, _) => {
                return Err(CliError::CommandArgumentError(
                    "Must choose either --yes or --no".to_string(),
                ))
            }
        };

        // TODO: Display details of proposal

        prompt_yes_with_override(
            &format!("Are you sure you want to vote {}", vote_str),
            self.prompt_options,
        )?;

        self.txn_options
            .submit_script_function(
                AccountAddress::ONE,
                "aptos_governance",
                "vote",
                vec![],
                vec![
                    bcs::to_bytes(&self.pool_address_args.pool_address)?,
                    bcs::to_bytes(&self.proposal_id)?,
                    bcs::to_bytes(&vote)?,
                ],
            )
            .await
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ProposalMetadata {
    title: String,
    description: String,
    source_code_url: String,
    discussion_url: String,
}

impl std::fmt::Display for ProposalMetadata {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Proposal:\n\tTitle:{}\n\tDescription:{}\n\tSource code URL:{}\n\tDiscussion URL:{}",
            self.title, self.description, self.source_code_url, self.discussion_url
        )
    }
}

fn compile_in_temp_dir(
    script_path: &Path,
    git_revision: &str,
    prompt_options: PromptOptions,
) -> CliTypedResult<(Vec<u8>, HashValue)> {
    // Make a temporary directory for compilation
    let temp_dir = TempDir::new().map_err(|err| {
        CliError::UnexpectedError(format!("Failed to create temporary directory {}", err))
    })?;

    // Initialize a move directory
    let package_dir = temp_dir.path();
    init_move_dir(
        package_dir,
        "Proposal",
        git_revision,
        BTreeMap::new(),
        prompt_options,
    )?;

    // Insert the new script
    let sources_dir = package_dir.join("sources");
    let new_script_path = if let Some(file_name) = script_path.file_name() {
        sources_dir.join(file_name)
    } else {
        // If for some reason we can't get the move file
        sources_dir.join("script.mv")
    };
    fs::copy(script_path, new_script_path.as_path()).map_err(|err| {
        CliError::IO(
            format!(
                "Failed to copy {} to {}",
                script_path.display(),
                new_script_path.display()
            ),
            err,
        )
    })?;

    // Compile the script
    compile_script(package_dir)
}

fn compile_script(package_dir: &Path) -> CliTypedResult<(Vec<u8>, HashValue)> {
    let build_config = BuildConfig {
        additional_named_addresses: BTreeMap::new(),
        generate_abis: false,
        generate_docs: false,
        ..Default::default()
    };

    let compiled_package = compile_move(build_config, package_dir)?;
    let scripts_count = compiled_package.scripts().count();

    if scripts_count != 1 {
        return Err(CliError::UnexpectedError(format!(
            "Only one script can be prepared a time. Make sure one and only one script file \
                is included in the Move package. Found {} scripts.",
            scripts_count
        )));
    }

    let script = *compiled_package
        .scripts()
        .collect::<Vec<_>>()
        .get(0)
        .unwrap();

    match script.unit {
        CompiledUnitEnum::Script(ref s) => {
            let mut bytes = vec![];

            s.script
                .serialize(&mut bytes)
                .map_err(|err| CliError::UnexpectedError(format!("Unexpected error: {}", err)))?;
            let hash = HashValue::sha3_256_of(bytes.as_slice());
            Ok((bytes, hash))
        }
        CompiledUnitEnum::Module(_) => Err(CliError::UnexpectedError(
            "You can only execute a script, a module is not supported.".to_string(),
        )),
    }
}

/// Execute a proposal that has passed voting requirements
#[derive(Parser)]
pub struct ExecuteProposal {
    /// Path to the Move script for the proposal
    #[clap(long, parse(from_os_str))]
    pub script_path: PathBuf,

    /// Git hash or branch of the framework in aptos core
    #[clap(long)]
    pub framework_git_rev: String,

    /// Arguments combined with their type separated by spaces.
    ///
    /// Supported types [u8, u64, u128, bool, hex, string, address]
    ///
    /// Example: `address:0x1 bool:true u8:0`
    #[clap(long, multiple_values = true)]
    pub(crate) args: Vec<ArgWithType>,

    /// TypeTag arguments separated by spaces.
    ///
    /// Example: `u8 u64 u128 bool address vector true false signer`
    #[clap(long, multiple_values = true)]
    pub(crate) type_args: Vec<MoveType>,

    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,

    #[clap(flatten)]
    pub(crate) prompt_options: PromptOptions,
}

impl TryFrom<&ArgWithType> for TransactionArgument {
    type Error = CliError;

    fn try_from(arg: &ArgWithType) -> Result<Self, Self::Error> {
        let txn_arg = match arg._ty {
            FunctionArgType::Address => TransactionArgument::Address(
                bcs::from_bytes(&arg.arg)
                    .map_err(|err| CliError::UnableToParse("address", err.to_string()))?,
            ),
            FunctionArgType::Bool => TransactionArgument::Bool(
                bcs::from_bytes(&arg.arg)
                    .map_err(|err| CliError::UnableToParse("bool", err.to_string()))?,
            ),
            FunctionArgType::Hex => TransactionArgument::U8Vector(
                bcs::from_bytes(&arg.arg)
                    .map_err(|err| CliError::UnableToParse("hex", err.to_string()))?,
            ),
            FunctionArgType::String => TransactionArgument::U8Vector(
                bcs::from_bytes(&arg.arg)
                    .map_err(|err| CliError::UnableToParse("string", err.to_string()))?,
            ),
            FunctionArgType::U128 => TransactionArgument::U128(
                bcs::from_bytes(&arg.arg)
                    .map_err(|err| CliError::UnableToParse("u128", err.to_string()))?,
            ),
            FunctionArgType::U64 => TransactionArgument::U64(
                bcs::from_bytes(&arg.arg)
                    .map_err(|err| CliError::UnableToParse("u64", err.to_string()))?,
            ),
            FunctionArgType::U8 => TransactionArgument::U8(
                bcs::from_bytes(&arg.arg)
                    .map_err(|err| CliError::UnableToParse("u8", err.to_string()))?,
            ),
        };

        Ok(txn_arg)
    }
}

#[async_trait]
impl CliCommand<TransactionSummary> for ExecuteProposal {
    fn command_name(&self) -> &'static str {
        "ExecuteProposal"
    }

    async fn execute(mut self) -> CliTypedResult<TransactionSummary> {
        // TODO: Make one object so this code isn't copied
        // Check script file
        let script_path = self.script_path.as_path();
        if self.script_path.exists() {
            return Err(CliError::CommandArgumentError(format!(
                "{} does not exist",
                script_path.display()
            )));
        } else if self.script_path.is_dir() {
            return Err(CliError::CommandArgumentError(format!(
                "{} is a directory",
                script_path.display()
            )));
        }

        // Compile script
        let (bytecode, _script_hash) =
            compile_in_temp_dir(script_path, &self.framework_git_rev, self.prompt_options)?;

        // TODO: Check hash so we don't do a failed roundtrip?

        // TODO: Clean these up to be common with the run function in move
        let args = self
            .args
            .iter()
            .map(|arg_with_type| arg_with_type.try_into())
            .collect::<Result<Vec<TransactionArgument>, CliError>>()?;

        let mut type_args: Vec<TypeTag> = Vec::new();

        // These TypeArgs are used for generics
        for type_arg in self.type_args.iter().cloned() {
            let type_tag = TypeTag::try_from(type_arg)
                .map_err(|err| CliError::UnableToParse("--type-args", err.to_string()))?;
            type_args.push(type_tag)
        }

        let txn = TransactionPayload::Script(Script::new(bytecode, type_args, args));

        self.txn_options
            .submit_transaction(txn)
            .await
            .map(TransactionSummary::from)
    }
}
<commit_msg>[aptos-cli] Considate duplicated script compilation code<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::common::types::{
    CliError, CliTypedResult, PoolAddressArgs, PromptOptions, TransactionOptions,
    TransactionSummary,
};
use crate::common::utils::prompt_yes_with_override;
use crate::move_tool::{compile_move, init_move_dir, ArgWithType, FunctionArgType};
use crate::{CliCommand, CliResult};
use aptos_crypto::HashValue;
use aptos_rest_client::{aptos_api_types::MoveType, Transaction};
use aptos_types::{
    account_address::AccountAddress,
    transaction::{Script, TransactionPayload},
};
use async_trait::async_trait;
use clap::Parser;
use move_deps::{
    move_compiler::compiled_unit::CompiledUnitEnum,
    move_core_types::{language_storage::TypeTag, transaction_argument::TransactionArgument},
    move_package::BuildConfig,
};
use reqwest::Url;
use serde::Deserialize;
use serde::Serialize;
use std::path::Path;
use std::{
    collections::BTreeMap,
    convert::{TryFrom, TryInto},
    fmt::Formatter,
    fs,
    path::PathBuf,
};
use tempfile::TempDir;

/// Tool for on-chain governance
///
/// This tool allows voters that have stake to vote the ability to
/// propose changes to the chain, as well as vote and execute these
/// proposals.
#[derive(Parser)]
pub enum GovernanceTool {
    Propose(SubmitProposal),
    Vote(SubmitVote),
    ExecuteProposal(ExecuteProposal),
}

impl GovernanceTool {
    pub async fn execute(self) -> CliResult {
        use GovernanceTool::*;
        match self {
            Propose(tool) => tool.execute_serialized().await,
            Vote(tool) => tool.execute_serialized().await,
            ExecuteProposal(tool) => tool.execute_serialized().await,
        }
    }
}

/// Submit proposal to other validators to be proposed on
#[derive(Parser)]
pub struct SubmitProposal {
    /// Code location of the script to be voted on
    #[clap(long)]
    pub(crate) metadata_url: Url,

    #[clap(flatten)]
    pub(crate) prompt_options: PromptOptions,
    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
    #[clap(flatten)]
    pub(crate) pool_address_args: PoolAddressArgs,
    #[clap(flatten)]
    pub(crate) compile_proposal_args: CompileProposalArgs,
}

#[async_trait]
impl CliCommand<Transaction> for SubmitProposal {
    fn command_name(&self) -> &'static str {
        "SubmitProposal"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        let (_bytecode, script_hash) = self.compile_proposal_args.compile()?;

        // Validate the proposal metadata
        let (metadata, metadata_hash) = get_metadata(self.metadata_url.clone()).await?;

        println!(
            "{}\n\tMetadata Hash: {}\n\tScript Hash: {}",
            metadata, metadata_hash, script_hash
        );
        prompt_yes_with_override("Do you want to submit this proposal?", self.prompt_options)?;

        self.txn_options
            .submit_script_function(
                AccountAddress::ONE,
                "aptos_governance",
                "create_proposal",
                vec![],
                vec![
                    bcs::to_bytes(&self.pool_address_args.pool_address)?,
                    bcs::to_bytes(&script_hash)?,
                    bcs::to_bytes(&self.metadata_url.to_string())?,
                    bcs::to_bytes(&metadata_hash.to_hex())?,
                ],
            )
            .await
    }
}

/// Retrieve metadata and validate it
async fn get_metadata(metadata_url: Url) -> CliTypedResult<(ProposalMetadata, HashValue)> {
    let client = reqwest::ClientBuilder::default()
        .tls_built_in_root_certs(true)
        .build()
        .map_err(|err| CliError::UnexpectedError(format!("Failed to build HTTP client {}", err)))?;
    let bytes = client
        .get(metadata_url.clone())
        .send()
        .await
        .map_err(|err| {
            CliError::CommandArgumentError(format!(
                "Failed to fetch metadata url {}: {}",
                metadata_url, err
            ))
        })?
        .bytes()
        .await
        .map_err(|err| {
            CliError::CommandArgumentError(format!(
                "Failed to fetch metadata url {}: {}",
                metadata_url, err
            ))
        })?;
    let metadata: ProposalMetadata = serde_json::from_slice(&bytes).map_err(|err| {
        CliError::CommandArgumentError(format!("Metadata is not in a proper JSON format: {}", err))
    })?;
    Url::parse(&metadata.source_code_url).map_err(|err| {
        CliError::CommandArgumentError(format!(
            "Source code URL {} is invalid {}",
            metadata.source_code_url, err
        ))
    })?;
    Url::parse(&metadata.discussion_url).map_err(|err| {
        CliError::CommandArgumentError(format!(
            "Discussion URL {} is invalid {}",
            metadata.discussion_url, err
        ))
    })?;
    let metadata_hash = HashValue::sha3_256_of(&bytes);
    Ok((metadata, metadata_hash))
}

/// Submit a vote on a current proposal
#[derive(Parser)]
pub struct SubmitVote {
    /// Id of proposal to vote on
    #[clap(long)]
    pub(crate) proposal_id: u64,

    /// Vote yes on the proposal
    #[clap(long, group = "vote")]
    pub(crate) yes: bool,

    /// Vote no on the proposal
    #[clap(long, group = "vote")]
    pub(crate) no: bool,

    #[clap(flatten)]
    pub(crate) prompt_options: PromptOptions,
    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
    #[clap(flatten)]
    pub(crate) pool_address_args: PoolAddressArgs,
}

#[async_trait]
impl CliCommand<Transaction> for SubmitVote {
    fn command_name(&self) -> &'static str {
        "SubmitVote"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        let (vote_str, vote) = match (self.yes, self.no) {
            (true, false) => ("Yes", true),
            (false, true) => ("No", false),
            (_, _) => {
                return Err(CliError::CommandArgumentError(
                    "Must choose either --yes or --no".to_string(),
                ))
            }
        };

        // TODO: Display details of proposal

        prompt_yes_with_override(
            &format!("Are you sure you want to vote {}", vote_str),
            self.prompt_options,
        )?;

        self.txn_options
            .submit_script_function(
                AccountAddress::ONE,
                "aptos_governance",
                "vote",
                vec![],
                vec![
                    bcs::to_bytes(&self.pool_address_args.pool_address)?,
                    bcs::to_bytes(&self.proposal_id)?,
                    bcs::to_bytes(&vote)?,
                ],
            )
            .await
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ProposalMetadata {
    title: String,
    description: String,
    source_code_url: String,
    discussion_url: String,
}

impl std::fmt::Display for ProposalMetadata {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Proposal:\n\tTitle:{}\n\tDescription:{}\n\tSource code URL:{}\n\tDiscussion URL:{}",
            self.title, self.description, self.source_code_url, self.discussion_url
        )
    }
}

fn compile_in_temp_dir(
    script_path: &Path,
    git_revision: &str,
    prompt_options: PromptOptions,
) -> CliTypedResult<(Vec<u8>, HashValue)> {
    // Make a temporary directory for compilation
    let temp_dir = TempDir::new().map_err(|err| {
        CliError::UnexpectedError(format!("Failed to create temporary directory {}", err))
    })?;

    // Initialize a move directory
    let package_dir = temp_dir.path();
    init_move_dir(
        package_dir,
        "Proposal",
        git_revision,
        BTreeMap::new(),
        prompt_options,
    )?;

    // Insert the new script
    let sources_dir = package_dir.join("sources");
    let new_script_path = if let Some(file_name) = script_path.file_name() {
        sources_dir.join(file_name)
    } else {
        // If for some reason we can't get the move file
        sources_dir.join("script.move")
    };
    fs::copy(script_path, new_script_path.as_path()).map_err(|err| {
        CliError::IO(
            format!(
                "Failed to copy {} to {}",
                script_path.display(),
                new_script_path.display()
            ),
            err,
        )
    })?;

    // Compile the script
    compile_script(package_dir)
}

fn compile_script(package_dir: &Path) -> CliTypedResult<(Vec<u8>, HashValue)> {
    let build_config = BuildConfig {
        additional_named_addresses: BTreeMap::new(),
        generate_abis: false,
        generate_docs: false,
        ..Default::default()
    };

    let compiled_package = compile_move(build_config, package_dir)?;
    let scripts_count = compiled_package.scripts().count();

    if scripts_count != 1 {
        return Err(CliError::UnexpectedError(format!(
            "Only one script can be prepared a time. Make sure one and only one script file \
                is included in the Move package. Found {} scripts.",
            scripts_count
        )));
    }

    let script = *compiled_package
        .scripts()
        .collect::<Vec<_>>()
        .get(0)
        .unwrap();

    match script.unit {
        CompiledUnitEnum::Script(ref s) => {
            let mut bytes = vec![];

            s.script
                .serialize(&mut bytes)
                .map_err(|err| CliError::UnexpectedError(format!("Unexpected error: {}", err)))?;
            let hash = HashValue::sha3_256_of(bytes.as_slice());
            Ok((bytes, hash))
        }
        CompiledUnitEnum::Module(_) => Err(CliError::UnexpectedError(
            "You can only execute a script, a module is not supported.".to_string(),
        )),
    }
}

/// Execute a proposal that has passed voting requirements
#[derive(Parser)]
pub struct ExecuteProposal {
    /// Arguments combined with their type separated by spaces.
    ///
    /// Supported types [u8, u64, u128, bool, hex, string, address]
    ///
    /// Example: `address:0x1 bool:true u8:0`
    #[clap(long, multiple_values = true)]
    pub(crate) args: Vec<ArgWithType>,

    /// TypeTag arguments separated by spaces.
    ///
    /// Example: `u8 u64 u128 bool address vector true false signer`
    #[clap(long, multiple_values = true)]
    pub(crate) type_args: Vec<MoveType>,

    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
    #[clap(flatten)]
    pub(crate) compile_proposal_args: CompileProposalArgs,
}

impl TryFrom<&ArgWithType> for TransactionArgument {
    type Error = CliError;

    fn try_from(arg: &ArgWithType) -> Result<Self, Self::Error> {
        let txn_arg = match arg._ty {
            FunctionArgType::Address => TransactionArgument::Address(
                bcs::from_bytes(&arg.arg)
                    .map_err(|err| CliError::UnableToParse("address", err.to_string()))?,
            ),
            FunctionArgType::Bool => TransactionArgument::Bool(
                bcs::from_bytes(&arg.arg)
                    .map_err(|err| CliError::UnableToParse("bool", err.to_string()))?,
            ),
            FunctionArgType::Hex => TransactionArgument::U8Vector(
                bcs::from_bytes(&arg.arg)
                    .map_err(|err| CliError::UnableToParse("hex", err.to_string()))?,
            ),
            FunctionArgType::String => TransactionArgument::U8Vector(
                bcs::from_bytes(&arg.arg)
                    .map_err(|err| CliError::UnableToParse("string", err.to_string()))?,
            ),
            FunctionArgType::U128 => TransactionArgument::U128(
                bcs::from_bytes(&arg.arg)
                    .map_err(|err| CliError::UnableToParse("u128", err.to_string()))?,
            ),
            FunctionArgType::U64 => TransactionArgument::U64(
                bcs::from_bytes(&arg.arg)
                    .map_err(|err| CliError::UnableToParse("u64", err.to_string()))?,
            ),
            FunctionArgType::U8 => TransactionArgument::U8(
                bcs::from_bytes(&arg.arg)
                    .map_err(|err| CliError::UnableToParse("u8", err.to_string()))?,
            ),
        };

        Ok(txn_arg)
    }
}

#[async_trait]
impl CliCommand<TransactionSummary> for ExecuteProposal {
    fn command_name(&self) -> &'static str {
        "ExecuteProposal"
    }

    async fn execute(mut self) -> CliTypedResult<TransactionSummary> {
        let (bytecode, _script_hash) = self.compile_proposal_args.compile()?;
        // TODO: Check hash so we don't do a failed roundtrip?

        // TODO: Clean these up to be common with the run function in move
        let args = self
            .args
            .iter()
            .map(|arg_with_type| arg_with_type.try_into())
            .collect::<Result<Vec<TransactionArgument>, CliError>>()?;

        let mut type_args: Vec<TypeTag> = Vec::new();

        // These TypeArgs are used for generics
        for type_arg in self.type_args.iter().cloned() {
            let type_tag = TypeTag::try_from(type_arg)
                .map_err(|err| CliError::UnableToParse("--type-args", err.to_string()))?;
            type_args.push(type_tag)
        }

        let txn = TransactionPayload::Script(Script::new(bytecode, type_args, args));

        self.txn_options
            .submit_transaction(txn)
            .await
            .map(TransactionSummary::from)
    }
}

/// Execute a proposal that has passed voting requirements
#[derive(Parser)]
pub struct CompileProposalArgs {
    /// Path to the Move script for the proposal
    #[clap(long, parse(from_os_str))]
    pub script_path: PathBuf,

    /// Git hash or branch of the framework in aptos core
    #[clap(long)]
    pub framework_git_rev: String,

    #[clap(flatten)]
    pub prompt_options: PromptOptions,
}

impl CompileProposalArgs {
    fn compile(self) -> CliTypedResult<(Vec<u8>, HashValue)> {
        // Check script file
        let script_path = self.script_path.as_path();
        if self.script_path.exists() {
            return Err(CliError::CommandArgumentError(format!(
                "{} does not exist",
                script_path.display()
            )));
        } else if self.script_path.is_dir() {
            return Err(CliError::CommandArgumentError(format!(
                "{} is a directory",
                script_path.display()
            )));
        }

        // Compile script
        compile_in_temp_dir(script_path, &self.framework_git_rev, self.prompt_options)
    }
}
<|endoftext|>


<commit_before>mod cargo_workspace;
mod json_project;
mod sysroot;

use std::{
    error::Error,
    fs::File,
    io::BufReader,
    path::{Path, PathBuf},
};

use ra_db::{CrateGraph, Edition, FileId};
use rustc_hash::FxHashMap;
use serde_json::from_reader;

pub use crate::{
    cargo_workspace::{CargoWorkspace, Package, Target, TargetKind},
    json_project::JsonProject,
    sysroot::Sysroot,
};

// FIXME use proper error enum
pub type Result<T> = ::std::result::Result<T, Box<dyn Error + Send + Sync>>;

#[derive(Debug, Clone)]
pub enum ProjectWorkspace {
    /// Project workspace was discovered by running `cargo metadata` and `rustc --print sysroot`.
    Cargo { cargo: CargoWorkspace, sysroot: Sysroot },
    /// Project workspace was manually specified using a `rust-project.json` file.
    Json { project: JsonProject },
}

/// `PackageRoot` describes a package root folder.
/// Which may be an external dependency, or a member of
/// the current workspace.
#[derive(Clone)]
pub struct PackageRoot {
    /// Path to the root folder
    path: PathBuf,
    /// Is a member of the current workspace
    is_member: bool,
}

impl PackageRoot {
    pub fn new(path: PathBuf, is_member: bool) -> PackageRoot {
        PackageRoot { path, is_member }
    }

    pub fn path(&self) -> &PathBuf {
        &self.path
    }

    pub fn is_member(&self) -> bool {
        self.is_member
    }
}

impl ProjectWorkspace {
    pub fn discover(path: &Path) -> Result<ProjectWorkspace> {
        ProjectWorkspace::discover_with_sysroot(path, true)
    }

    pub fn discover_with_sysroot(path: &Path, with_sysroot: bool) -> Result<ProjectWorkspace> {
        match find_rust_project_json(path) {
            Some(json_path) => {
                let file = File::open(json_path)?;
                let reader = BufReader::new(file);
                Ok(ProjectWorkspace::Json { project: from_reader(reader)? })
            }
            None => {
                let cargo_toml = find_cargo_toml(path)?;
                let cargo = CargoWorkspace::from_cargo_metadata(&cargo_toml)?;
                let sysroot =
                    if with_sysroot { Sysroot::discover(&cargo_toml)? } else { Sysroot::default() };
                Ok(ProjectWorkspace::Cargo { cargo, sysroot })
            }
        }
    }

    /// Returns the roots for the current `ProjectWorkspace`
    /// The return type contains the path and whether or not
    /// the root is a member of the current workspace
    pub fn to_roots(&self) -> Vec<PackageRoot> {
        match self {
            ProjectWorkspace::Json { project } => {
                let mut roots = Vec::with_capacity(project.roots.len());
                for root in &project.roots {
                    roots.push(PackageRoot::new(root.path.clone(), true));
                }
                roots
            }
            ProjectWorkspace::Cargo { cargo, sysroot } => {
                let mut roots = Vec::with_capacity(cargo.packages().len() + sysroot.crates().len());
                for pkg in cargo.packages() {
                    let root = pkg.root(&cargo).to_path_buf();
                    let member = pkg.is_member(&cargo);
                    roots.push(PackageRoot::new(root, member));
                }
                for krate in sysroot.crates() {
                    roots.push(PackageRoot::new(krate.root_dir(&sysroot).to_path_buf(), false))
                }
                roots
            }
        }
    }

    pub fn n_packages(&self) -> usize {
        match self {
            ProjectWorkspace::Json { project } => project.crates.len(),
            ProjectWorkspace::Cargo { cargo, sysroot } => {
                cargo.packages().len() + sysroot.crates().len()
            }
        }
    }

    pub fn to_crate_graph(&self, load: &mut dyn FnMut(&Path) -> Option<FileId>) -> CrateGraph {
        let mut crate_graph = CrateGraph::default();
        match self {
            ProjectWorkspace::Json { project } => {
                let mut crates = FxHashMap::default();
                for (id, krate) in project.crates.iter().enumerate() {
                    let crate_id = json_project::CrateId(id);
                    if let Some(file_id) = load(&krate.root_module) {
                        let edition = match krate.edition {
                            json_project::Edition::Edition2015 => Edition::Edition2015,
                            json_project::Edition::Edition2018 => Edition::Edition2018,
                        };
                        crates.insert(crate_id, crate_graph.add_crate_root(file_id, edition));
                    }
                }

                for (id, krate) in project.crates.iter().enumerate() {
                    for dep in &krate.deps {
                        let from_crate_id = json_project::CrateId(id);
                        let to_crate_id = dep.krate;
                        if let (Some(&from), Some(&to)) =
                            (crates.get(&from_crate_id), crates.get(&to_crate_id))
                        {
                            if let Err(_) = crate_graph.add_dep(from, dep.name.clone().into(), to) {
                                log::error!(
                                    "cyclic dependency {:?} -> {:?}",
                                    from_crate_id,
                                    to_crate_id
                                );
                            }
                        }
                    }
                }
            }
            ProjectWorkspace::Cargo { cargo, sysroot } => {
                let mut sysroot_crates = FxHashMap::default();
                for krate in sysroot.crates() {
                    if let Some(file_id) = load(krate.root(&sysroot)) {
                        sysroot_crates.insert(
                            krate,
                            crate_graph.add_crate_root(file_id, Edition::Edition2015),
                        );
                    }
                }
                for from in sysroot.crates() {
                    for to in from.deps(&sysroot) {
                        let name = to.name(&sysroot);
                        if let (Some(&from), Some(&to)) =
                            (sysroot_crates.get(&from), sysroot_crates.get(&to))
                        {
                            if let Err(_) = crate_graph.add_dep(from, name.into(), to) {
                                log::error!("cyclic dependency between sysroot crates")
                            }
                        }
                    }
                }

                let libstd = sysroot.std().and_then(|it| sysroot_crates.get(&it).copied());

                let mut pkg_to_lib_crate = FxHashMap::default();
                let mut pkg_crates = FxHashMap::default();
                // Next, create crates for each package, target pair
                for pkg in cargo.packages() {
                    let mut lib_tgt = None;
                    for tgt in pkg.targets(&cargo) {
                        let root = tgt.root(&cargo);
                        if let Some(file_id) = load(root) {
                            let edition = pkg.edition(&cargo);
                            let crate_id = crate_graph.add_crate_root(file_id, edition);
                            if tgt.kind(&cargo) == TargetKind::Lib {
                                lib_tgt = Some(crate_id);
                                pkg_to_lib_crate.insert(pkg, crate_id);
                            }
                            pkg_crates.entry(pkg).or_insert_with(Vec::new).push(crate_id);
                        }
                    }

                    // Set deps to the std and to the lib target of the current package
                    for &from in pkg_crates.get(&pkg).into_iter().flatten() {
                        if let Some(to) = lib_tgt {
                            if to != from {
                                if let Err(_) =
                                    crate_graph.add_dep(from, pkg.name(&cargo).into(), to)
                                {
                                    log::error!(
                                        "cyclic dependency between targets of {}",
                                        pkg.name(&cargo)
                                    )
                                }
                            }
                        }
                        if let Some(std) = libstd {
                            if let Err(_) = crate_graph.add_dep(from, "std".into(), std) {
                                log::error!("cyclic dependency on std for {}", pkg.name(&cargo))
                            }
                        }
                    }
                }

                // Now add a dep ednge from all targets of upstream to the lib
                // target of downstream.
                for pkg in cargo.packages() {
                    for dep in pkg.dependencies(&cargo) {
                        if let Some(&to) = pkg_to_lib_crate.get(&dep.pkg) {
                            for &from in pkg_crates.get(&pkg).into_iter().flatten() {
                                if let Err(_) =
                                    crate_graph.add_dep(from, dep.name.clone().into(), to)
                                {
                                    log::error!(
                                        "cyclic dependency {} -> {}",
                                        pkg.name(&cargo),
                                        dep.pkg.name(&cargo)
                                    )
                                }
                            }
                        }
                    }
                }
            }
        }
        crate_graph
    }

    pub fn workspace_root_for(&self, path: &Path) -> Option<&Path> {
        match self {
            ProjectWorkspace::Cargo { cargo, .. } => {
                Some(cargo.workspace_root.as_ref()).filter(|root| path.starts_with(root))
            }
            ProjectWorkspace::Json { project: JsonProject { roots, .. } } => roots
                .iter()
                .find(|root| path.starts_with(&root.path))
                .map(|root| root.path.as_ref()),
        }
    }
}

fn find_rust_project_json(path: &Path) -> Option<PathBuf> {
    if path.ends_with("rust-project.json") {
        return Some(path.to_path_buf());
    }

    let mut curr = Some(path);
    while let Some(path) = curr {
        let candidate = path.join("rust-project.json");
        if candidate.exists() {
            return Some(candidate);
        }
        curr = path.parent();
    }

    None
}

fn find_cargo_toml(path: &Path) -> Result<PathBuf> {
    if path.ends_with("Cargo.toml") {
        return Ok(path.to_path_buf());
    }
    let mut curr = Some(path);
    while let Some(path) = curr {
        let candidate = path.join("Cargo.toml");
        if candidate.exists() {
            return Ok(candidate);
        }
        curr = path.parent();
    }
    Err(format!("can't find Cargo.toml at {}", path.display()))?
}
<commit_msg>close #1763<commit_after>mod cargo_workspace;
mod json_project;
mod sysroot;

use std::{
    error::Error,
    fs::File,
    io::BufReader,
    path::{Path, PathBuf},
};

use ra_db::{CrateGraph, Edition, FileId};
use rustc_hash::FxHashMap;
use serde_json::from_reader;

pub use crate::{
    cargo_workspace::{CargoWorkspace, Package, Target, TargetKind},
    json_project::JsonProject,
    sysroot::Sysroot,
};

// FIXME use proper error enum
pub type Result<T> = ::std::result::Result<T, Box<dyn Error + Send + Sync>>;

#[derive(Debug, Clone)]
pub enum ProjectWorkspace {
    /// Project workspace was discovered by running `cargo metadata` and `rustc --print sysroot`.
    Cargo { cargo: CargoWorkspace, sysroot: Sysroot },
    /// Project workspace was manually specified using a `rust-project.json` file.
    Json { project: JsonProject },
}

/// `PackageRoot` describes a package root folder.
/// Which may be an external dependency, or a member of
/// the current workspace.
#[derive(Clone)]
pub struct PackageRoot {
    /// Path to the root folder
    path: PathBuf,
    /// Is a member of the current workspace
    is_member: bool,
}

impl PackageRoot {
    pub fn new(path: PathBuf, is_member: bool) -> PackageRoot {
        PackageRoot { path, is_member }
    }

    pub fn path(&self) -> &PathBuf {
        &self.path
    }

    pub fn is_member(&self) -> bool {
        self.is_member
    }
}

impl ProjectWorkspace {
    pub fn discover(path: &Path) -> Result<ProjectWorkspace> {
        ProjectWorkspace::discover_with_sysroot(path, true)
    }

    pub fn discover_with_sysroot(path: &Path, with_sysroot: bool) -> Result<ProjectWorkspace> {
        match find_rust_project_json(path) {
            Some(json_path) => {
                let file = File::open(json_path)?;
                let reader = BufReader::new(file);
                Ok(ProjectWorkspace::Json { project: from_reader(reader)? })
            }
            None => {
                let cargo_toml = find_cargo_toml(path)?;
                let cargo = CargoWorkspace::from_cargo_metadata(&cargo_toml)?;
                let sysroot =
                    if with_sysroot { Sysroot::discover(&cargo_toml)? } else { Sysroot::default() };
                Ok(ProjectWorkspace::Cargo { cargo, sysroot })
            }
        }
    }

    /// Returns the roots for the current `ProjectWorkspace`
    /// The return type contains the path and whether or not
    /// the root is a member of the current workspace
    pub fn to_roots(&self) -> Vec<PackageRoot> {
        match self {
            ProjectWorkspace::Json { project } => {
                let mut roots = Vec::with_capacity(project.roots.len());
                for root in &project.roots {
                    roots.push(PackageRoot::new(root.path.clone(), true));
                }
                roots
            }
            ProjectWorkspace::Cargo { cargo, sysroot } => {
                let mut roots = Vec::with_capacity(cargo.packages().len() + sysroot.crates().len());
                for pkg in cargo.packages() {
                    let root = pkg.root(&cargo).to_path_buf();
                    let member = pkg.is_member(&cargo);
                    roots.push(PackageRoot::new(root, member));
                }
                for krate in sysroot.crates() {
                    roots.push(PackageRoot::new(krate.root_dir(&sysroot).to_path_buf(), false))
                }
                roots
            }
        }
    }

    pub fn n_packages(&self) -> usize {
        match self {
            ProjectWorkspace::Json { project } => project.crates.len(),
            ProjectWorkspace::Cargo { cargo, sysroot } => {
                cargo.packages().len() + sysroot.crates().len()
            }
        }
    }

    pub fn to_crate_graph(&self, load: &mut dyn FnMut(&Path) -> Option<FileId>) -> CrateGraph {
        let mut crate_graph = CrateGraph::default();
        match self {
            ProjectWorkspace::Json { project } => {
                let mut crates = FxHashMap::default();
                for (id, krate) in project.crates.iter().enumerate() {
                    let crate_id = json_project::CrateId(id);
                    if let Some(file_id) = load(&krate.root_module) {
                        let edition = match krate.edition {
                            json_project::Edition::Edition2015 => Edition::Edition2015,
                            json_project::Edition::Edition2018 => Edition::Edition2018,
                        };
                        crates.insert(crate_id, crate_graph.add_crate_root(file_id, edition));
                    }
                }

                for (id, krate) in project.crates.iter().enumerate() {
                    for dep in &krate.deps {
                        let from_crate_id = json_project::CrateId(id);
                        let to_crate_id = dep.krate;
                        if let (Some(&from), Some(&to)) =
                            (crates.get(&from_crate_id), crates.get(&to_crate_id))
                        {
                            if let Err(_) = crate_graph.add_dep(from, dep.name.clone().into(), to) {
                                log::error!(
                                    "cyclic dependency {:?} -> {:?}",
                                    from_crate_id,
                                    to_crate_id
                                );
                            }
                        }
                    }
                }
            }
            ProjectWorkspace::Cargo { cargo, sysroot } => {
                let mut sysroot_crates = FxHashMap::default();
                for krate in sysroot.crates() {
                    if let Some(file_id) = load(krate.root(&sysroot)) {
                        sysroot_crates.insert(
                            krate,
                            crate_graph.add_crate_root(file_id, Edition::Edition2018),
                        );
                    }
                }
                for from in sysroot.crates() {
                    for to in from.deps(&sysroot) {
                        let name = to.name(&sysroot);
                        if let (Some(&from), Some(&to)) =
                            (sysroot_crates.get(&from), sysroot_crates.get(&to))
                        {
                            if let Err(_) = crate_graph.add_dep(from, name.into(), to) {
                                log::error!("cyclic dependency between sysroot crates")
                            }
                        }
                    }
                }

                let libstd = sysroot.std().and_then(|it| sysroot_crates.get(&it).copied());

                let mut pkg_to_lib_crate = FxHashMap::default();
                let mut pkg_crates = FxHashMap::default();
                // Next, create crates for each package, target pair
                for pkg in cargo.packages() {
                    let mut lib_tgt = None;
                    for tgt in pkg.targets(&cargo) {
                        let root = tgt.root(&cargo);
                        if let Some(file_id) = load(root) {
                            let edition = pkg.edition(&cargo);
                            let crate_id = crate_graph.add_crate_root(file_id, edition);
                            if tgt.kind(&cargo) == TargetKind::Lib {
                                lib_tgt = Some(crate_id);
                                pkg_to_lib_crate.insert(pkg, crate_id);
                            }
                            pkg_crates.entry(pkg).or_insert_with(Vec::new).push(crate_id);
                        }
                    }

                    // Set deps to the std and to the lib target of the current package
                    for &from in pkg_crates.get(&pkg).into_iter().flatten() {
                        if let Some(to) = lib_tgt {
                            if to != from {
                                if let Err(_) =
                                    crate_graph.add_dep(from, pkg.name(&cargo).into(), to)
                                {
                                    log::error!(
                                        "cyclic dependency between targets of {}",
                                        pkg.name(&cargo)
                                    )
                                }
                            }
                        }
                        if let Some(std) = libstd {
                            if let Err(_) = crate_graph.add_dep(from, "std".into(), std) {
                                log::error!("cyclic dependency on std for {}", pkg.name(&cargo))
                            }
                        }
                    }
                }

                // Now add a dep ednge from all targets of upstream to the lib
                // target of downstream.
                for pkg in cargo.packages() {
                    for dep in pkg.dependencies(&cargo) {
                        if let Some(&to) = pkg_to_lib_crate.get(&dep.pkg) {
                            for &from in pkg_crates.get(&pkg).into_iter().flatten() {
                                if let Err(_) =
                                    crate_graph.add_dep(from, dep.name.clone().into(), to)
                                {
                                    log::error!(
                                        "cyclic dependency {} -> {}",
                                        pkg.name(&cargo),
                                        dep.pkg.name(&cargo)
                                    )
                                }
                            }
                        }
                    }
                }
            }
        }
        crate_graph
    }

    pub fn workspace_root_for(&self, path: &Path) -> Option<&Path> {
        match self {
            ProjectWorkspace::Cargo { cargo, .. } => {
                Some(cargo.workspace_root.as_ref()).filter(|root| path.starts_with(root))
            }
            ProjectWorkspace::Json { project: JsonProject { roots, .. } } => roots
                .iter()
                .find(|root| path.starts_with(&root.path))
                .map(|root| root.path.as_ref()),
        }
    }
}

fn find_rust_project_json(path: &Path) -> Option<PathBuf> {
    if path.ends_with("rust-project.json") {
        return Some(path.to_path_buf());
    }

    let mut curr = Some(path);
    while let Some(path) = curr {
        let candidate = path.join("rust-project.json");
        if candidate.exists() {
            return Some(candidate);
        }
        curr = path.parent();
    }

    None
}

fn find_cargo_toml(path: &Path) -> Result<PathBuf> {
    if path.ends_with("Cargo.toml") {
        return Ok(path.to_path_buf());
    }
    let mut curr = Some(path);
    while let Some(path) = curr {
        let candidate = path.join("Cargo.toml");
        if candidate.exists() {
            return Ok(candidate);
        }
        curr = path.parent();
    }
    Err(format!("can't find Cargo.toml at {}", path.display()))?
}
<|endoftext|>