Amend GasSchedule Updator CLI for easy scaling of minimum Gas unit price.

aptos-move/aptos-gas-schedule-updator/Cargo.toml

@@ -21,6 +21,8 @@ move-model = { workspace = true }
 
 anyhow = { workspace = true }
 bcs = { workspace = true }
+serde = { workspace = true }
+serde_json = { workspace = true }
 clap = { workspace = true }
 
 [dev-dependencies]

aptos-move/aptos-gas-schedule-updator/src/change_set.rs

@@ -0,0 +1,75 @@
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+
+pub type Entries = HashMap<String, u64>;
+
+/// A set of changes to be applied to a gas schedule.
+/// Additions are new entries to be added to the gas schedule.
+/// Deletions are entries to be removed from the gas schedule.
+/// Mutations are entries to be updated in the gas schedule.
+#[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)]
+pub struct GasScheduleChangeSet {
+    additions: Entries,
+    deletions: Entries,
+    mutations: Entries,
+}
+
+impl GasScheduleChangeSet {
+
+    /// Deserialize a GasScheduleChangeSet from a JSON string.
+    pub fn from_json_string(json_str: String) -> anyhow::Result<Self> {
+        serde_json::from_str(&json_str).map_err(|e| anyhow::anyhow!(e))
+    }
+
+    pub fn deletions(&self) -> &Entries {
+        &self.deletions
+    }
+
+    pub fn additions(&self) -> &Entries {
+        &self.additions
+    }
+
+    pub fn mutations(&self) -> &Entries {
+        &self.mutations
+    }
+
+    /// Returns true if the change set is empty.
+    pub fn is_empty(&self) -> bool {
+        self.additions.is_empty() && self.deletions.is_empty() && self.mutations.is_empty()
+    }
+
+    /// Returns true if the change set contains any additions or deletions.
+    /// This indicates that a feature version bump is required.
+    /// A mutation alone does not require a feature version bump.
+    // Check `aptos-core/aptos-move/aptos-gas-schedule/src/ver.rs` for more context.
+    pub fn should_bump_feature_version(&self) -> bool {
+        !self.additions.is_empty() || !self.deletions.is_empty()
+    }
+}
+
+#[test]
+// Test deserialization of GasScheduleChangeSet from JSON string.
+fn test_deserialize_change_set() {
+    let json = r#"{
+        "additions": {
+             "foo": 123,
+             "bar": 100999
+        },
+        "deletions": {
+            "bar": 456
+        },
+        "mutations": {
+            "foo": 789
+        }
+    }"#
+    .to_string();
+
+    let change_set = GasScheduleChangeSet::from_json_string(json).unwrap();
+
+    assert_eq!(change_set.additions.len(), 2);
+    assert_eq!(change_set.deletions.len(), 1);
+    assert_eq!(change_set.mutations.len(), 1);
+    assert_eq!(change_set.additions.get("foo").unwrap(), &123u64);
+    assert_eq!(change_set.additions.get("bar").unwrap(), &100999u64);
+    assert_eq!(change_set.deletions.get("foo").unwrap(), &789u64);
+}

aptos-move/aptos-gas-schedule-updator/src/lib.rs

@@ -7,17 +7,23 @@
 //! The generated proposal includes a comment section, listing the contents of the
 //! gas schedule in a human readable format.
 
-use anyhow::Result;
+mod change_set;
+
+use crate::change_set::GasScheduleChangeSet;
+use anyhow::{anyhow, Result};
 use aptos_gas_schedule::{
     AptosGasParameters, InitialGasSchedule, ToOnChainGasSchedule, LATEST_GAS_FEATURE_VERSION,
 };
 use aptos_package_builder::PackageBuilder;
 use aptos_types::on_chain_config::GasScheduleV2;
-use clap::Parser;
+use clap::{Args, Parser};
 use move_core_types::account_address::AccountAddress;
 use move_model::{code_writer::CodeWriter, emit, emitln, model::Loc};
+use std::fs;
 use std::path::{Path, PathBuf};
 
+const DEFAULT_GAS_SCHEDULE_SCRIPT_UPDATE_PATH: &str = "./proposals";
+
 fn generate_blob(writer: &CodeWriter, data: &[u8]) {
     emitln!(writer, "vector[");
     writer.indent();
@@ -37,7 +43,7 @@ fn generate_blob(writer: &CodeWriter, data: &[u8]) {
 }
 
 fn generate_script(gas_schedule: &GasScheduleV2) -> Result<String> {
-    let gas_schedule_blob = bcs::to_bytes(gas_schedule).unwrap();
+    let gas_schedule_blob = bcs::to_bytes(gas_schedule)?;
 
     assert!(gas_schedule_blob.len() < 65536);
 
@@ -103,16 +109,187 @@ fn aptos_framework_path() -> PathBuf {
     )
 }
 
+/// Command line interface for the gas schedule update proposal generation tool.
+/// It supports two modes:
+/// 1. Generate a new gas schedule from the current hardcoded values.
+/// 2. Update an existing gas schedule with a change set.
+/// The generated proposal is written to a Move package in the specified output directory.
+/// If no output directory is specified, it defaults to `./proposals`.
+/// The generated package contains a single Move script `update_gas_schedule.move`.
+/// This script can be submitted as a governance proposal to update the on-chain gas schedule.
+/// The script includes a comment section listing the contents of the gas schedule in a human readable format.
+#[derive(Parser, Debug)]
+pub enum GasScheduleGenerator {
+    /// Generate a new gas schedule from the current hardcoded values.
+    /// Optionally specify the feature version of the gas schedule.
+    /// If not specified, it defaults to the latest feature version.
+    GenerateNew(GenerateNewSchedule),
+    /// Update an existing gas schedule with a change set.
+    /// The change set is specified in a JSON file.
+    /// The current gas schedule is also specified in a JSON file.
+    /// The change set can include additions, deletions, and mutations of gas parameters.
+    /// If the change set includes any additions or deletions, the feature version of the gas
+    /// schedule is bumped by 1.
+    UpdateSchedule(UpdateSchedule),
+}
+
 /// Command line arguments to the gas schedule update proposal generation tool.
-#[derive(Debug, Parser)]
-pub struct GenArgs {
-    #[clap(short, long)]
+#[derive(Debug, Args)]
+pub struct GenerateNewSchedule {
+    /// Path to file to write the output script.
+    /// If not specified, it defaults to `./proposals`.
+    #[clap(short, long, help = "Path to file to write the output script")]
     pub output: Option<String>,
 
-    #[clap(short, long)]
+    /// Feature version of the gas schedule to generate.
+    /// If not specified, it defaults to the latest feature version.
+    #[clap(short, long, help = "Feature version of the GasSchedule generated")]
     pub gas_feature_version: Option<u64>,
 }
 
+impl GenerateNewSchedule {
+    pub fn execute(self) -> Result<()> {
+        let feature_version = self
+            .gas_feature_version
+            .unwrap_or(LATEST_GAS_FEATURE_VERSION);
+
+        let gas_schedule = current_gas_schedule(feature_version);
+
+        generate_update_proposal(
+            &gas_schedule,
+            self.output
+                .unwrap_or_else(|| DEFAULT_GAS_SCHEDULE_SCRIPT_UPDATE_PATH.to_string()),
+        )
+    }
+}
+
+#[derive(Parser, Debug)]
+pub struct UpdateSchedule {
+    /// Path to file to write the output script.
+    /// If not specified, it defaults to `./proposals`.
+    #[clap(short, long, help = "Path to file to write the output script")]
+    pub output: Option<String>,
+
+    /// Path to JSON file containing the current GasScheduleV2 to update.
+    /// The JSON file should be in the format produced by the `to_json_string` method
+    /// of the `GasScheduleV2` struct.
+    #[clap(
+        short,
+        long,
+        help = "Path to JSON file containing the GasScheduleV2 to update"
+    )]
+    pub current_schedule_path: String,
+
+    /// Path to JSON file containing the change set to apply to the current GasScheduleV2.
+    /// The JSON file should be in the format produced by the `to_json_string` method
+    /// of the `GasScheduleChangeSet` struct.
+    #[clap(
+        short,
+        long,
+        help = "Path to JSON file containing change set to the GasScheduleV2"
+    )]
+    pub change_set_path: String,
+}
+
+impl UpdateSchedule {
+    pub fn execute(self) -> Result<()> {
+        let change_set_json_str = fs::read_to_string(self.change_set_path)?;
+        let change_set = GasScheduleChangeSet::from_json_string(change_set_json_str)?;
+
+        if change_set.is_empty() {
+            return Err(anyhow::anyhow!("The change set is empty"));
+        }
+
+        let current_schedule_json_str = fs::read_to_string(self.current_schedule_path)?;
+        let mut current_schedule = GasScheduleV2::from_json_string(current_schedule_json_str)?;
+
+        // Apply additions to the gas schedule
+        for (name, value) in change_set.additions().iter() {
+            if current_schedule
+                .entries
+                .iter()
+                .any(|entry| entry.0 == *name && entry.1 == *value)
+            {
+                return Err(anyhow!(
+                    "Addition entry ({}, {}) is already found in GasSchedule",
+                    name,
+                    value
+                ));
+            }
+
+            current_schedule.entries.push((name.clone(), *value));
+        }
+
+        // Apply deletions to existing entries in the gas schedule
+        for (name, value) in change_set.deletions().iter() {
+            if let Some(position) = current_schedule
+                .entries
+                .iter()
+                .position(|entry| entry.0 == *name && entry.1 == *value)
+            {
+                current_schedule.entries.remove(position);
+            } else {
+                return Err(anyhow!(
+                    "Deletion entry ({}, {}) not found in GasSchedule",
+                    name,
+                    value
+                ));
+            }
+        }
+
+        // Apply mutations to existing entries in the gas schedule
+        for (name, value) in change_set.mutations().iter() {
+            if let Some(position) = current_schedule
+                .entries
+                .iter()
+                .position(|entry| entry.0 == *name )
+            {
+                // Remove old entry and insert new entry
+                current_schedule.entries.remove(position);
+                current_schedule.entries.push((name.clone(), *value));
+            } else {
+                return Err(anyhow!(
+                    "Mutation entry ({}, {}) not found in GasSchedule",
+                    name,
+                    value
+                ));
+            }
+        }
+
+
+        // Sort the entries by name to ensure deterministic order
+        current_schedule.entries.sort_by(|a, b| a.0.cmp(&b.0));
+        // Ensure no duplicate entries exist
+        let mut seen = std::collections::HashSet::new();
+        for (name, _) in &current_schedule.entries {
+            if !seen.insert(name) {
+                return Err(anyhow!("Duplicate entry ({}) found in GasSchedule", name));
+            }
+        }
+
+        // Bump the feature version if we are adding or removing params
+        if change_set.should_bump_feature_version() {
+            let new_feature_version = current_schedule
+                .feature_version
+                .checked_add(1)
+                .ok_or_else(|| anyhow!("Overflow when bumping feature version"))?;
+            current_schedule.feature_version = new_feature_version;
+        }
+
+        println!(
+            "Updated gas schedule to feature version {}",
+            current_schedule.feature_version
+        );
+
+        // Generate the update proposal script
+        generate_update_proposal(
+            &current_schedule,
+            self.output
+                .unwrap_or_else(|| DEFAULT_GAS_SCHEDULE_SCRIPT_UPDATE_PATH.to_string()),
+        )
+    }
+}
+
 /// Constructs the current gas schedule in on-chain format.
 pub fn current_gas_schedule(feature_version: u64) -> GasScheduleV2 {
     GasScheduleV2 {
@@ -122,27 +299,23 @@ pub fn current_gas_schedule(feature_version: u64) -> GasScheduleV2 {
 }
 
 /// Entrypoint for the update proposal generation tool.
-pub fn generate_update_proposal(args: &GenArgs) -> Result<()> {
+pub fn generate_update_proposal(gas_schedule: &GasScheduleV2, output_path: String) -> Result<()> {
     let mut pack = PackageBuilder::new("GasScheduleUpdate");
 
-    let feature_version = args
-        .gas_feature_version
-        .unwrap_or(LATEST_GAS_FEATURE_VERSION);
-
-    pack.add_source(
-        "update_gas_schedule.move",
-        &generate_script(&current_gas_schedule(feature_version))?,
-    );
+    pack.add_source("update_gas_schedule.move", &generate_script(gas_schedule)?);
     // TODO: use relative path here
     pack.add_local_dep("SupraFramework", &aptos_framework_path().to_string_lossy());
 
-    pack.write_to_disk(args.output.as_deref().unwrap_or("./proposal"))?;
+    pack.write_to_disk(PathBuf::from(output_path))?;
 
     Ok(())
 }
 
-#[test]
-fn verify_tool() {
-    use clap::CommandFactory;
-    GenArgs::command().debug_assert()
+impl GasScheduleGenerator {
+    pub fn execute(self) -> Result<()> {
+        match self {
+            GasScheduleGenerator::GenerateNew(args) => args.execute(),
+            GasScheduleGenerator::UpdateSchedule(args) => args.execute(),
+        }
+    }
 }

aptos-move/aptos-gas-schedule-updator/src/main.rs

@@ -2,11 +2,9 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use anyhow::Result;
-use aptos_gas_schedule_updator::{generate_update_proposal, GenArgs};
+use aptos_gas_schedule_updator::GasScheduleGenerator;
 use clap::Parser;
 
 fn main() -> Result<()> {
-    let args = GenArgs::parse();
-
-    generate_update_proposal(&args)
+    GasScheduleGenerator::parse().execute()
 }

aptos-move/aptos-gas-schedule-updator/tests/gen_tests.rs

@@ -2,16 +2,17 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use aptos_framework::{BuildOptions, BuiltPackage};
-use aptos_gas_schedule_updator::{generate_update_proposal, GenArgs};
+use aptos_gas_schedule_updator::GenerateNewSchedule;
 
 #[test]
 fn can_generate_and_build_update_proposal() {
     let output_dir = tempfile::tempdir().unwrap();
 
-    generate_update_proposal(&GenArgs {
+    GenerateNewSchedule {
         gas_feature_version: None,
         output: Some(output_dir.path().to_string_lossy().to_string()),
-    })
+    }
+    .execute()
     .unwrap();
 
     BuiltPackage::build(output_dir.path().to_path_buf(), BuildOptions::default()).unwrap();